Ax Flow

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

const llm = ai({ name: 'openai', apiKey: process.env.OPENAI_APIKEY! });

const wf = flow<{ userInput: string }, { responseText: string }>()
  .node('testNode', 'userInput:string -> responseText:string')
  .execute('testNode', (state) => ({ userInput: state.userInput }))
  .returns((state) => ({ responseText: state.testNodeResult.responseText }));

const result = await wf.forward(llm, { userInput: 'Hello world' });
console.log(result.responseText);

// Basic
const wf = flow();

// With options
const wf = flow({ autoParallel: false });

// Typed
const wf = flow<InputType, OutputType>();

// Typed with options
const wf = flow<InputType, OutputType>({ autoParallel: true, batchSize: 5 });

// Initial state: { userInput: 'Hello' }
flow.execute('processor', (state) => ({ input: state.userInput }));
// State: { userInput: 'Hello', processorResult: { output: '...' } }

flow.execute('analyzer', (state) => ({ text: state.processorResult.output }));
// State: { ..., analyzerResult: { sentiment: '...', confidence: 0.8 } }

// String signature (creates AxGen automatically)
flow.node('processor', 'input:string -> output:string');

// Multiple outputs
flow.node('analyzer', 'text:string -> sentiment:string, confidence:number');

// Array outputs
flow.node('extractor', 'documentText:string -> entities:string[]');

// Short alias
flow.n('processor', 'input:string -> output:string');

flow
  .node('triage', 'ticketText:string -> ticketClass:class "bug, billing, question", severityScore:number(min 1, max 5)')
  .node('draft', 'ticketText:string, ticketClass:string, severityScore:number -> replyText:string(max 400)')
  .node('audit', 'replyText:string -> approved:boolean, flaggedSpans:object{ spanText:string, reasonNote:string }[]');

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

// Chain-of-thought reasoning
flow.nx('reasoner', 'question:string -> answer:string', {
  prependOutputs: [
    { name: 'reasoning', type: f.internal(f.string('Step-by-step reasoning')) },
  ],
});

// Add confidence scoring
flow.nx('analyzer', 'input:string -> result:string', {
  appendOutputs: [{ name: 'confidence', type: f.number('Confidence 0-1') }],
});

// Add optional context input
flow.nx('processor', 'query:string -> response:string', {
  appendInputs: [{ name: 'context', type: f.optional(f.string('Extra context')) }],
});

flow.execute('summarizer', (state) => ({ documentText: state.document }));

// With AI override (use a different model for this node)
flow.execute('processor', (state) => ({ input: state.data }), { ai: alternativeAI });

// Sync
flow.map((state) => ({ ...state, upperText: state.rawText.toUpperCase() }));

// Async
flow.map(async (state) => {
  const data = await fetchFromAPI(state.query);
  return { ...state, enrichedData: data };
});

// Parallel async transforms
flow.map([
  async (state) => ({ ...state, result1: await api1(state.data) }),
  async (state) => ({ ...state, result2: await api2(state.data) }),
], { parallel: true });

const wf = flow<{ input: string }>()
  .map((state) => ({ ...state, upper: state.input.toUpperCase(), len: state.input.length }))
  .returns((state) => ({ upper: state.upper, isLong: state.len > 20 }));

// Result is typed as { upper: string; isLong: boolean }
const result = await wf.forward(llm, { input: 'test' });

const wf = flow<{ input: string }, { finalResult: string }>()
  .node('step1', 'input:string -> intermediate:string')
  .node('step2', 'intermediate:string -> output:string')
  .execute('step1', (state) => ({ input: state.input }))
  .execute('step2', (state) => ({ intermediate: state.step1Result.intermediate }))
  .returns((state) => ({ finalResult: state.step2Result.output }));

const wf = flow<{ text: string }, { combined: string }>()
  .node('sentimentAnalyzer', 'text:string -> sentiment:string')
  .node('topicExtractor', 'text:string -> topics:string[]')
  .node('entityRecognizer', 'text:string -> entities:string[]')
  // These three run in parallel (all depend only on state.text)
  .execute('sentimentAnalyzer', (state) => ({ text: state.text }))
  .execute('topicExtractor', (state) => ({ text: state.text }))
  .execute('entityRecognizer', (state) => ({ text: state.text }))
  // This waits for all three
  .returns((state) => ({
    combined: JSON.stringify({
      sentiment: state.sentimentAnalyzerResult.sentiment,
      topics: state.topicExtractorResult.topics,
      entities: state.entityRecognizerResult.entities,
    }),
  }));

// Inspect execution plan
const plan = wf.getExecutionPlan();
console.log(plan.parallelGroups, plan.maxParallelism);

const wf = flow({ autoParallel: false });
// or per execution:
await wf.forward(llm, input, { autoParallel: false });

const wf = flow<{ query: string; expertMode: boolean }, { response: string }>()
  .node('simple', 'query:string -> response:string')
  .node('expert', 'query:string -> response:string')
  .branch((state) => state.expertMode)
    .when(true)
      .execute('expert', (state) => ({ query: state.query }))
    .when(false)
      .execute('simple', (state) => ({ query: state.query }))
  .merge()
  .returns((state) => ({
    response: state.expertResult?.response ?? state.simpleResult?.response,
  }));

const wf = flow<{ content: string }, { finalContent: string }>()
  .node('processor', 'content:string -> processedContent:string')
  .node('qualityChecker', 'content:string -> qualityScore:number')
  .map((state) => ({ currentContent: state.content, iteration: 0, qualityScore: 0 }))
  .while((state) => state.iteration < 3 && state.qualityScore < 0.8)
    .map((state) => ({ ...state, iteration: state.iteration + 1 }))
    .execute('processor', (state) => ({ content: state.currentContent }))
    .execute('qualityChecker', (state) => ({
      content: state.processorResult.processedContent,
    }))
    .map((state) => ({
      ...state,
      currentContent: state.processorResult.processedContent,
      qualityScore: state.qualityCheckerResult.qualityScore,
    }))
  .endWhile()
  .returns((state) => ({ finalContent: state.currentContent }));

const wf = flow<{ prompt: string }, { result: string }>()
  .node('gen', 'prompt:string -> result:string, quality:number')
  .map((state) => ({ ...state, tries: 0 }))
  .label('retry')
    .map((state) => ({ ...state, tries: state.tries + 1 }))
    .execute('gen', (state) => ({ prompt: state.prompt }))
  .feedback((state) => state.genResult.quality < 0.9 && state.tries < 3, 'retry')
  .returns((state) => ({ result: state.genResult.result }));

flow
  .parallel([
    (sub) => sub.execute('analyzer1', (state) => ({ text: state.input })),
    (sub) => sub.execute('analyzer2', (state) => ({ text: state.input })),
    (sub) => sub.execute('analyzer3', (state) => ({ text: state.input })),
  ])
  .merge('combinedResults', (r1, r2, r3) => ({
    a1: r1.analyzer1Result.analysis,
    a2: r2.analyzer2Result.analysis,
    a3: r3.analyzer3Result.analysis,
  }));

const wf = flow<{ items: string[] }, { processed: string[] }>({ batchSize: 3 })
  .derive('processed', 'items', (item, index) => `processed-${item}-${index}`, {
    batchSize: 2,
  });

const fast = ai({ name: 'openai', apiKey: '...', config: { model: 'gpt-5.4-mini' } });
const smart = ai({ name: 'anthropic', apiKey: '...' });

const wf = flow<{ text: string }, { out: string }>()
  .node('draft', 'text:string -> out:string')
  .node('refine', 'text:string -> out:string')
  .execute('draft', (state) => ({ text: state.text }), { ai: fast })
  .execute('refine', (state) => ({ text: state.draftResult.out }), { ai: smart })
  .returns((state) => ({ out: state.refineResult.out }));

const wf = flow<{ userQuestion: string }, { responseText: string }>()
  .node('qa', 'userQuestion:string -> responseText:string')
  .execute('qa', (state) => ({ userQuestion: state.userQuestion }))
  .returns((state) => ({ responseText: state.qaResult.responseText }))
  .description('Question Answerer', 'Answers user questions concisely.');

const fn = wf.toFunction();
// fn.name, fn.parameters (JSON Schema), fn.func

import { ai, flow } from '@ax-llm/ax';
import { context, trace } from '@opentelemetry/api';

const tracer = trace.getTracer('axflow');
const llm = ai({ name: 'openai', apiKey: '...' });

const wf = flow<{ userQuestion: string }>()
  .node('summarizer', 'documentText:string -> summaryText:string')
  .execute('summarizer', (s) => ({ documentText: s.userQuestion }))
  .returns((s) => ({ answer: s.summarizerResult.summaryText }));

const result = await wf.forward(llm, { userQuestion: 'hi' }, {
  tracer,
  traceContext: context.active(),
});

import { axGlobals } from '@ax-llm/ax';
import { metrics } from '@opentelemetry/api';

axGlobals.tracer = tracer;
axGlobals.meter = metrics.getMeter('axflow');
axGlobals.rateLimiter = async (next, info) => next();

const result = await wf.forward(llm, { userQuestion: 'hi' });

const wf = flow<{ input: string }>()
  .node('summarizer', 'text:string -> summary:string')
  .node('classifier', 'text:string -> category:string');

// Discover program IDs
console.log(wf.namedPrograms());
// [{ id: 'root.summarizer', ... }, { id: 'root.classifier', ... }]

// Set demos (TypeScript catches typos)
wf.setDemos([{ programId: 'root.summarizer', traces: [] }]);

// Apply optimization
wf.applyOptimization(optimizedProgram);

const log = wf.getChatLog();
for (const entry of log) {
  console.log(entry.name, entry.model);
}

try {
  const result = await wf.forward(llm, input);
} catch (error) {
  console.error('Flow execution failed:', error);
}

const wf = flow({ mcp: [inventory], ucp: [merchant] })
  .node('lookup', lookupProgram)
  .node('checkout', checkoutProgram, { mcpInheritance: ['merchant'] });

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

const wf = flow<{ documentText: string }, { finalReport: string }>(`
flowchart TD
  %%ax summarize: documentText:string -> summaryText:string(max 500)
  %%ax check: summaryText:string -> verdict:class "pass, fail", note?:string
  %%ax format: summaryText:string, note?:string -> finalReport:string

  summarize[Summarize document] --> check{verdict}
  check -->|pass| format
  check -->|fail, max 3| summarize
`);

const { finalReport } = await wf.forward(llm, { documentText });
console.log(String(wf)); // render back to the same dialect