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, 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));