import { AxAIOpenAIModel, ai, flow } from '@ax-llm/ax';

const apiKey = process.env.OPENAI_API_KEY ?? process.env.OPENAI_APIKEY;
if (!apiKey) {
  throw new Error('Set OPENAI_API_KEY or OPENAI_APIKEY to run this example.');
}

const llm = ai({
  name: 'openai',
  apiKey,
  config: {
    model: AxAIOpenAIModel.GPT54Mini,
    temperature: 0,
  },
});

const workflow = flow<{ documentText: string }>()
  .description(
    'TypeScript Sequential Flow',
    'Runs a two-step Ax flow against OpenAI.'
  )
  .node('summarizer', 'documentText:string -> summaryText:string')
  .node(
    'classifier',
    'textContent:string -> priority:class "high, normal, low"'
  )
  .execute('summarizer', (state) => ({ documentText: state.documentText }))
  .execute('classifier', (state) => ({
    textContent: state.summarizerResult.summaryText,
  }))
  .returns((state) => ({
    summary: state.summarizerResult.summaryText as string,
    priority: state.classifierResult.priority as string,
  }));

const result = await workflow.forward(llm, {
  documentText:
    'Ax gives developers typed signatures, provider clients, agents, flows, tracing, and optimization so LLM features can be built as ordinary programs.',
});

console.log(JSON.stringify(result, null, 2));

import { AxAIOpenAIModel, ai, flow } from '@ax-llm/ax';

const apiKey = process.env.OPENAI_API_KEY ?? process.env.OPENAI_APIKEY;
if (!apiKey) {
  throw new Error('Set OPENAI_API_KEY or OPENAI_APIKEY to run this example.');
}

const llm = ai({
  name: 'openai',
  apiKey,
  config: {
    model: AxAIOpenAIModel.GPT54Mini,
    temperature: 0,
  },
});

const workflow = flow<{ requestText: string }>()
  .description(
    'TypeScript Branching Flow',
    'Routes a classification through follow-up flow logic backed by OpenAI.'
  )
  .node(
    'classifier',
    'requestText:string -> route:class "support, sales, engineering"'
  )
  .node('responder', 'requestText:string, route:string -> responseText:string')
  .execute('classifier', (state) => ({ requestText: state.requestText }))
  .execute('responder', (state) => ({
    requestText: state.requestText,
    route: state.classifierResult.route,
  }))
  .returns((state) => ({
    route: state.classifierResult.route as string,
    responseText: state.responderResult.responseText as string,
  }));

const result = await workflow.forward(llm, {
  requestText: 'A customer says checkout is down for their enterprise account.',
});

console.log(JSON.stringify(result, null, 2));

import { AxAIOpenAIModel, ai, flow } from '@ax-llm/ax';

const apiKey = process.env.OPENAI_API_KEY ?? process.env.OPENAI_APIKEY;
if (!apiKey) {
  throw new Error('Set OPENAI_API_KEY or OPENAI_APIKEY to run this example.');
}

const llm = ai({
  name: 'openai',
  apiKey,
  config: {
    model: AxAIOpenAIModel.GPT54Mini,
    temperature: 0,
  },
});

const workflow = flow<{ topicText: string }>()
  .description(
    'TypeScript Parallel Flow',
    'Research and audience analysis run independently before the join step.'
  )
  .node('research', 'topicText:string -> factList:string[]')
  .node('audience', 'topicText:string -> audienceAngle:string')
  .node(
    'join',
    'factList:string[], audienceAngle:string -> briefText:string(max 500)'
  )
  .execute('research', (state) => ({ topicText: state.topicText }))
  .execute('audience', (state) => ({ topicText: state.topicText }))
  .execute('join', (state) => ({
    factList: state.researchResult.factList,
    audienceAngle: state.audienceResult.audienceAngle,
  }))
  .returns((state) => ({ briefText: state.joinResult.briefText }));

const result = await workflow.forward(llm, {
  topicText:
    'Why typed contracts make multi-step LLM systems easier to maintain',
});

console.log(JSON.stringify(result, null, 2));

import { AxAIOpenAIModel, ai, ax, flow, fn } from '@ax-llm/ax';

const calls = { research: 0, inventory: 0 };
const research = fn('lookupResearch')
  .description('Get the research reference; call once.')
  .execution('background')
  .handler(async () => {
    calls.research++;
    await new Promise((resolve) => setTimeout(resolve, 2000));
    return 'RESEARCH-314';
  })
  .build();
const inventory = fn('lookupInventory')
  .description('Get the inventory reference; call once.')
  .execution('background')
  .handler(async () => {
    calls.inventory++;
    await new Promise((resolve) => setTimeout(resolve, 500));
    return 'STOCK-271';
  })
  .build();
const workflow = flow<{ question: string }>()
  .node('research', ax('question -> answer', { functions: [research] }))
  .node('inventory', ax('question -> answer', { functions: [inventory] }))
  .execute('research', () => ({
    question: 'Call lookupResearch and report its exact reference.',
  }))
  .execute('inventory', () => ({
    question: 'Call lookupInventory and report its exact reference.',
  }))
  .returns((state) => ({
    references: {
      research: state.researchResult.answer,
      inventory: state.inventoryResult.answer,
    },
  }));

const result = await workflow.forward(
  ai({
    name: 'openai',
    apiKey: process.env.OPENAI_API_KEY ?? process.env.OPENAI_APIKEY,
    config: { model: AxAIOpenAIModel.GPT6Astra, maxTokens: 1500 },
  }),
  { question: 'Get both references' },
  {
    thinkingTokenBudget: 'low',
    serviceTier: 'standard',
    abortSignal: AbortSignal.timeout(120_000),
  }
);
if (calls.research !== 1 || calls.inventory !== 1)
  throw new Error('Expected each lookup to execute once');
if (
  !result.references.research.includes('RESEARCH-314') ||
  !result.references.inventory.includes('STOCK-271')
)
  throw new Error('Flow omitted a tool result');
if (
  result.references.research.includes('STOCK-271') ||
  result.references.inventory.includes('RESEARCH-314')
)
  throw new Error('Parallel conversation results leaked across nodes');
console.log(result.references);

import { AxAIOpenAIModel, ai, flow } from '@ax-llm/ax';

const apiKey = process.env.OPENAI_API_KEY ?? process.env.OPENAI_APIKEY;
if (!apiKey) {
  throw new Error('Set OPENAI_API_KEY or OPENAI_APIKEY to run this example.');
}

const llm = ai({
  name: 'openai',
  apiKey,
  config: {
    model: AxAIOpenAIModel.GPT54Mini,
    temperature: 0,
  },
});

const workflow = flow<{ topic: string }>()
  .description(
    'TypeScript Composed Flow',
    'Composes multiple typed programs into one OpenAI-backed flow.'
  )
  .node('outline', 'topic:string -> outline:string[]')
  .node('brief', 'topic:string, outline:string[] -> brief:string')
  .execute('outline', (state) => ({ topic: state.topic }))
  .execute('brief', (state) => ({
    topic: state.topic,
    outline: state.outlineResult.outline,
  }))
  .returns((state) => ({
    outline: state.outlineResult.outline as string[],
    brief: state.briefResult.brief as string,
  }));

const result = await workflow.forward(llm, {
  topic:
    'How Ax moves from typed generation to agents, flows, and optimization',
});

console.log(JSON.stringify(result, null, 2));

import { AxAIOpenAIModel, ai, flow } from '@ax-llm/ax';

const apiKey = process.env.OPENAI_API_KEY ?? process.env.OPENAI_APIKEY;
if (!apiKey) {
  throw new Error('Set OPENAI_API_KEY or OPENAI_APIKEY to run this example.');
}

const llm = ai({
  name: 'openai',
  apiKey,
  config: {
    model: AxAIOpenAIModel.GPT54Mini,
    temperature: 0,
  },
});

const workflow = flow<{ topicText: string }>()
  .description(
    'TypeScript Refinement Flow',
    'A linear draft, critique, and revision pipeline.'
  )
  .node('draft', 'topicText:string -> draftText:string(max 500)')
  .node('critique', 'draftText:string -> critiqueText:string(max 250)')
  .node(
    'revise',
    'draftText:string, critiqueText:string -> revisedText:string(max 800)'
  )
  .execute('draft', (state) => ({ topicText: state.topicText }))
  .execute('critique', (state) => ({
    draftText: state.draftResult.draftText,
  }))
  .execute('revise', (state) => ({
    draftText: state.draftResult.draftText,
    critiqueText: state.critiqueResult.critiqueText,
  }))
  .returns((state) => ({ revisedText: state.reviseResult.revisedText }));

const result = await workflow.forward(llm, {
  topicText: 'Explain automatic flow parallelism to a backend engineer.',
});

console.log(JSON.stringify(result, null, 2));

import { AxAIOpenAIModel, ai, ax, flow, fn, runControl } from '@ax-llm/ax';

const control = runControl();
const appliedPaths: string[] = [];
control.onEvent((event) => {
  if (event.type === 'applied') appliedPaths.push(event.path);
});
let baselineCalls = 0;
let reviewCalls = 0;
const baseline = fn('lookupBaseline')
  .description('Get the baseline reference; call once.')
  .execution('background')
  .handler(() => {
    baselineCalls++;
    return 'BASE-101';
  })
  .build();
const review = fn('lookupVerification')
  .description('Get the verification reference; call once.')
  .execution('background')
  .handler(() => {
    reviewCalls++;
    return 'CHECK-202';
  })
  .build();
const workflow = flow<{ question: string }>()
  .node('baseline', ax('question -> answer', { functions: [baseline] }))
  .node('review', ax('question -> answer', { functions: [review] }))
  .execute('baseline', (state) => ({ question: state.question }))
  .execute('review', (state) => {
    control.steer(
      'Include both exact references and the word REVIEWED in the final answer.',
      { target: 'root/review' }
    );
    control.setThinkingTokenBudget('medium', { target: 'root/review' });
    return {
      question: `Use lookupVerification to review this baseline: ${state.baselineResult.answer}`,
    };
  })
  .returns((state) => ({
    baseline: state.baselineResult.answer,
    review: state.reviewResult.answer,
  }));
const result = await workflow.forward(
  ai({
    name: 'openai',
    apiKey: process.env.OPENAI_API_KEY ?? process.env.OPENAI_APIKEY,
    config: { model: AxAIOpenAIModel.GPT6Astra, maxTokens: 2000 },
  }),
  { question: 'Call lookupBaseline and report its exact reference.' },
  {
    control,
    thinkingTokenBudget: 'low',
    serviceTier: 'standard',
    abortSignal: AbortSignal.timeout(120_000),
  }
);
if (baselineCalls !== 1 || reviewCalls !== 1)
  throw new Error('A completed node was rerun or a tool was skipped');
if (
  appliedPaths.length !== 2 ||
  appliedPaths.some((path) => path !== 'root/review')
)
  throw new Error('Updates were applied outside the review scope');
if (
  !['BASE-101', 'CHECK-202', 'REVIEWED'].every((text) =>
    result.review.includes(text)
  ) ||
  result.baseline.includes('REVIEWED')
)
  throw new Error('Targeted instructions were not isolated or incorporated');
console.log(result);

import { AxAIOpenAIModel, ai, ax, flow, fn, runControl } from '@ax-llm/ax';

const control = runControl();
let started = 0;
let cancelled = 0;
let completed = false;
control.onEvent((event) => {
  if (event.type === 'completed' && event.path === 'root') completed = true;
});
const pendingLookup = () =>
  fn('lookup')
    .description('Start the required lookup. Call exactly once.')
    .execution('background')
    .handler(async (_args, extra) => {
      const signal = extra?.abortSignal;
      if (!signal) throw new Error('Tool did not receive cancellation context');
      signal.throwIfAborted();
      await new Promise<never>((_resolve, reject) => {
        signal.addEventListener(
          'abort',
          () => {
            cancelled++;
            reject(signal.reason);
          },
          { once: true }
        );
        started++;
        // Simulate a user pressing Stop after both lookups have begun.
        if (started === 2) control.abort();
      });
      return 'unreachable';
    })
    .build();
const workflow = flow<{ question: string }>()
  .node('left', ax('question -> answer', { functions: [pendingLookup()] }))
  .node('right', ax('question -> answer', { functions: [pendingLookup()] }))
  .execute('left', (state) => ({ question: state.question }))
  .execute('right', (state) => ({ question: state.question }))
  .returns((state) => ({
    answers: [state.leftResult.answer, state.rightResult.answer],
  }));

let rejected = false;
try {
  await workflow.forward(
    ai({
      name: 'openai',
      apiKey: process.env.OPENAI_API_KEY ?? process.env.OPENAI_APIKEY,
      config: { model: AxAIOpenAIModel.GPT6Astra, maxTokens: 1500 },
    }),
    { question: 'Call lookup to obtain the required reference.' },
    {
      control,
      thinkingTokenBudget: 'low',
      serviceTier: 'standard',
      abortSignal: AbortSignal.timeout(120_000),
    }
  );
} catch (error) {
  if (!control.signal.aborted) throw error;
  rejected = true;
}
if (!rejected || started !== 2 || cancelled !== 2 || completed)
  throw new Error('Cancellation did not stop both pending branches');
console.log(
  'Cancelled both pending lookups; the flow did not report completion.'
);