import { AxAIOpenAIModel, ai, ax } 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 program = ax('question:string -> answer:string');
const result = await program.forward(llm, {
  question:
    'In one sentence, explain Ax as a language-agnostic LLM programming library.',
});

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

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

const apiKey = process.env.TYPESAFE_API_KEY ?? process.env.TYPESAFE_APIKEY;
if (!apiKey)
  throw new Error(
    'Set TYPESAFE_API_KEY or TYPESAFE_APIKEY to run this example.'
  );

const model = ai({ name: 'typesafe', apiKey, trueThreshold: 0.9 });
const triage = ax(
  `ticket:string -> urgent:boolean(
    true "Customers cannot complete a core task",
    false "A routine request or minor inconvenience"
  ) "Does this need immediate attention?",
  team:class "support, billing, engineering"(
    support "Product usage questions",
    billing "Individual invoice or charge disputes",
    engineering "Broken functionality or service outages"
  ) "Which team should investigate?"`
);

const decision = await triage.forward(model, {
  ticket:
    'Checkout is returning errors for every customer. Payments cannot complete.',
});

console.log(decision);
console.log(triage.getChatLog().at(-1)?.providerMetadata?.typesafe?.answers);

import '../../meta-muse.js';

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

const llm = ai({
  name: 'openai',
  apiKey: process.env.OPENAI_API_KEY ?? process.env.OPENAI_APIKEY,
  config: { model: AxAIOpenAIModel.GPT6Astra },
});
const answer = await ax('question:string -> answer:string').forward(
  llm,
  { question: 'Why do leaves change color in autumn?' },
  { thinkingTokenBudget: 'low', serviceTier: 'standard' }
);
console.log(answer);

import { AxAIOpenAIModel, ai, ax } 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 program = ax(
  'ticket:string -> priority:class "high, normal, low", summary:string, labels:string[]'
);
const result = await program.forward(llm, {
  ticket:
    'Checkout has failed for enterprise customers since 09:00. Support wants a concise summary and tags.',
});

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

import { AxAIOpenAIModel, ai, ax, f, s } 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 bookingSignature = f()
  .input('requestText', f.string('Booking request').min(10).max(500))
  .input('contactEmail', f.string('Contact email').email())
  .output('partySize', f.number('Guests').min(1).max(12))
  .output(
    'bookingCode',
    f
      .string('Three letters, a dash, and four digits')
      .regex('^[A-Z]{3}-\\d{4}$', 'Must look like ABC-1234')
  )
  .output(
    'guestProfile',
    f.object({
      fullName: f.string('Primary guest').min(2),
      dietaryNotes: f.string('Dietary requirements').optional(),
    })
  )
  .build();

const extendedStringSignature = s(
  'requestText:string -> booking:object{ bookingCode:string(pattern "^[A-Z]{3}-\\\\d{4}$" "ABC-1234"), partySize:number(min 1, max 12) }'
);

const result = await ax(bookingSignature).forward(llm, {
  requestText: 'Book dinner for four people under the name Ada Lovelace.',
  contactEmail: 'ada@example.com',
});

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

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

const typesafeKey = process.env.TYPESAFE_API_KEY ?? process.env.TYPESAFE_APIKEY;
const openaiKey = process.env.OPENAI_API_KEY ?? process.env.OPENAI_APIKEY;
if (!typesafeKey || !openaiKey) {
  throw new Error(
    'Set TYPESAFE_API_KEY (or TYPESAFE_APIKEY) and OPENAI_API_KEY (or OPENAI_APIKEY) to run this example.'
  );
}

const typesafe = ai({
  name: 'typesafe',
  apiKey: typesafeKey,
  trueThreshold: 0.9,
});
const writer = ai({
  name: 'openai',
  apiKey: openaiKey,
  config: { model: AxAIOpenAIModel.GPT56Luna },
});

const ticket =
  'Checkout is returning errors for every customer. Payments cannot complete.';
const triage = ax(
  'ticket:string -> urgent:boolean, team:class "support, billing, engineering"'
);
const decision = await triage.forward(typesafe, { ticket });

const summarize = ax(
  'ticket:string, urgent:boolean, team:string -> summary:string "Brief internal summary based on the supplied triage decisions"'
);
const { summary } = await summarize.forward(writer, { ticket, ...decision });

console.log({ ...decision, summary });

import { AxAIOpenAIModel, type AxUsageEvent, ai, axGlobals } 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 events: Readonly<AxUsageEvent>[] = [];
axGlobals.onUsage = (event) => {
  // In production, enqueue this synchronously and persist it out of band.
  events.push(event);
};

const llm = ai({
  name: 'openai',
  apiKey,
  config: { model: AxAIOpenAIModel.GPT54Mini, temperature: 0 },
  options: {
    usageContext: {
      tenantId: 'tenant-42',
      feature: 'support-chat',
      attributes: { environment: 'example' },
    },
  },
});

try {
  await llm.chat(
    {
      chatPrompt: [{ role: 'user', content: 'Reply with one short greeting.' }],
    },
    {
      usageContext: {
        userId: 'user-7',
        requestId: crypto.randomUUID(),
      },
    }
  );
  console.log(JSON.stringify(events, null, 2));
} finally {
  axGlobals.onUsage = undefined;
}

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

const apiKey = process.env.TYPESAFE_API_KEY ?? process.env.TYPESAFE_APIKEY;
if (!apiKey)
  throw new Error(
    'Set TYPESAFE_API_KEY or TYPESAFE_APIKEY to run this example.'
  );

const client = typesafe({ apiKey });
const result = await client.systemOne({
  state: {
    ticket: {
      text: 'Checkout returns errors for every customer. Payments cannot complete.',
      affectedFeature: 'payments',
    },
    recentEvents: ['Checkout deployment completed', 'Payment errors increased'],
  },
  questions: {
    urgent: {
      type: 'noul',
      instructions: 'Does this require immediate incident response?',
      criteria: {
        true: {
          description: 'Customers cannot complete a core task',
          examples: ['Payments are unavailable'],
        },
        false: 'A routine request or minor inconvenience',
      },
    },
    team: {
      type: 'choice',
      instructions: 'Which team should investigate?',
      criteria: {
        support: 'Product usage questions',
        billing: 'Individual invoice or charge disputes',
        engineering: 'Broken functionality or service outages',
      },
    },
    severity: {
      type: 'score',
      instructions: 'How severe is the customer impact?',
      criteria: [
        'Minor inconvenience with a workaround',
        'A core feature is impaired for some customers',
        'A core feature is unavailable for all customers',
      ],
    },
  },
});

// These are application policies. Native Noul and Score results remain unchanged.
const escalate = result.answers.urgent.noul >= 0.9;
const severity = result.answers.severity.score; // Fractional rubric position, 0–2.
const models = await client.listModels();
console.log({
  escalate,
  severity,
  answers: result.answers,
  usage: result.usage,
});
console.log({ availableModels: models.map((model) => model.name) });

import assert from 'node:assert/strict';
import { AxAIOpenAIModel, ai, ax } from '@ax-llm/ax';

const typesafeKey = process.env.TYPESAFE_API_KEY ?? process.env.TYPESAFE_APIKEY;
const openaiKey = process.env.OPENAI_API_KEY ?? process.env.OPENAI_APIKEY;
if (!typesafeKey || !openaiKey)
  throw new Error('Set Typesafe and OpenAI API keys to run this example.');

const triage = ax(`
  ticket:string ->
  urgent:boolean(
    true "Customers cannot complete a core task",
    false "A routine request or minor inconvenience"
  ) "Does this need immediate attention?",
  team:class "support, billing, engineering"(
    support "Product usage questions",
    billing "Individual invoice or charge disputes",
    engineering "Broken functionality or service outages"
  ) "Which team should investigate?"
`);

const providers = [
  ai({ name: 'typesafe', apiKey: typesafeKey, trueThreshold: 0.9 }),
  ai({
    name: 'openai',
    apiKey: openaiKey,
    config: { model: AxAIOpenAIModel.GPT56Luna },
  }),
];

for (const model of providers) {
  const decision = await triage.forward(model, {
    ticket:
      'Checkout returns errors for every customer. Payments cannot complete.',
  });
  assert.equal(typeof decision.urgent, 'boolean');
  assert.ok(['support', 'billing', 'engineering'].includes(decision.team));
  console.log(model.getName(), decision);
}

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

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

const model = AxAIGoogleGeminiModel.Gemini38Flash;
const catalogModel = axGetSupportedAIModels()
  .find((provider) => provider.name === 'google-gemini')
  ?.models.find((candidate) => candidate.name === model);

if (!catalogModel?.capabilities.serviceTiers.includes('flex')) {
  throw new Error(`${model} does not advertise Gemini Flex support.`);
}

console.log(
  `Thinking levels: ${catalogModel.capabilities.thinkingLevels.join(', ')}`
);

const gemini = ai({
  name: 'google-gemini',
  apiKey,
  config: {
    model,
  },
});

const result = await gemini.chat(
  {
    chatPrompt: [
      {
        role: 'user',
        content:
          'Summarize why batching independent evaluation work saves time.',
      },
    ],
  },
  { stream: false, serviceTier: 'flex' }
);

if (result instanceof ReadableStream) {
  throw new Error('Expected a non-streaming Gemini response.');
}

console.log(result.results[0]?.content);
console.log(
  `Handled by: ${result.modelUsage?.tokens?.serviceTier ?? 'unknown'}`
);

import { AxAIBedrock, AxAIBedrockModel } from '@ax-llm/ax-ai-aws-bedrock';

const model = AxAIBedrockModel.ClaudeSonnet5;
const bedrock = new AxAIBedrock({
  region: process.env.AWS_REGION ?? 'us-east-2',
  fallbackRegions: ['us-west-2', 'us-east-1'],
  config: { model, maxTokens: 4096 },
});

const weather = {
  name: 'current_weather',
  description: 'Read the current weather for a city',
  parameters: {
    type: 'object',
    properties: {
      city: { type: 'string', description: 'City name' },
    },
    required: ['city'],
  },
};

const userMessage = {
  role: 'user' as const,
  content: 'What should I wear for today in Vancouver?',
};
const toolRequest = await bedrock.chat(
  {
    model,
    chatPrompt: [
      { role: 'system', content: 'Use tools for current facts.', cache: true },
      userMessage,
    ],
    functions: [weather],
    functionCall: 'required',
  },
  {
    contextCache: { ttlSeconds: 3600 },
    thinkingTokenBudget: 'medium',
  }
);

if (toolRequest instanceof ReadableStream) {
  throw new Error('Expected the tool request to be non-streaming.');
}

const assistant = toolRequest.results[0];
const call = assistant?.functionCalls?.[0];
if (!call) throw new Error('Claude did not request the weather tool.');

const toolResult = JSON.stringify({
  city: 'Vancouver',
  conditions: 'light rain',
  temperatureC: 14,
});

const stream = await bedrock.chat(
  {
    model,
    chatPrompt: [
      { role: 'system', content: 'Use tools for current facts.', cache: true },
      userMessage,
      {
        role: 'assistant',
        content: assistant.content,
        functionCalls: assistant.functionCalls,
        thoughtBlocks: assistant.thoughtBlocks,
      },
      { role: 'function', functionId: call.id, result: toolResult },
    ],
    functions: [weather],
  },
  {
    stream: true,
    contextCache: { ttlSeconds: 3600 },
    thinkingTokenBudget: 'medium',
  }
);

if (!(stream instanceof ReadableStream)) {
  throw new Error('Expected a streaming Bedrock response.');
}

const reader = stream.getReader();
while (true) {
  const { done, value } = await reader.read();
  if (done) break;
  process.stdout.write(value.results[0]?.content ?? '');
}
process.stdout.write('\n');

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

const control = runControl();
const llm = ai({
  name: 'openai',
  apiKey: process.env.OPENAI_API_KEY ?? process.env.OPENAI_APIKEY,
  config: { model: AxAIOpenAIModel.GPT6Astra, maxTokens: 2000 },
});
const lookup = fn('lookupCode')
  .description('Look up the required confirmation code')
  .execution('background')
  .handler(async () => {
    control.setThinkingTokenBudget('high');
    control.steer(
      'Include the confirmation code verbatim and keep the answer brief.'
    );
    await new Promise((resolve) => setTimeout(resolve, 2000));
    return 'AX-742';
  })
  .build();
const workflow = flow<{ question: string }>()
  .node('delivery', ax('question -> answer', { functions: [lookup] }))
  .execute('delivery', (state) => ({ question: state.question }))
  .returns((state) => ({ answer: state.deliveryResult.answer }));
const result = await workflow.forward(
  llm,
  { question: 'Call lookupCode and report the confirmation code.' },
  {
    control,
    thinkingTokenBudget: 'low',
    serviceTier: 'standard',
    abortSignal: AbortSignal.timeout(120_000),
  }
);
if (!result.answer.includes('AX-742'))
  throw new Error('Flow did not incorporate the tool result');
console.log(result.answer);

import { AxAIOpenAIModel, ai, ax, runControl } from '@ax-llm/ax';
import WebSocket from 'ws';

const control = runControl();
const llm = ai({
  name: 'openai',
  apiKey: process.env.OPENAI_API_KEY ?? process.env.OPENAI_APIKEY,
  config: { model: AxAIOpenAIModel.GPT6Astra, maxTokens: 2500 },
  options: { webSocket: WebSocket },
});
let steered = false;
control.onEvent((event) => {
  if (event.type === 'model.output' && !steered) {
    steered = true;
    control.steer(
      'Change the plan: use exactly three short bullet points and include the phrase Small launch.'
    );
  }
  if (event.type === 'applied')
    console.log(`Steering applied: ${event.timing}`);
});
let answer = '';
let version = -1;
for await (const chunk of ax('question -> answer').streamingForward(
  llm,
  {
    question:
      'Write a detailed ten-step launch plan for a task tracking application.',
  },
  {
    control,
    thinkingTokenBudget: 'low',
    serviceTier: 'standard',
    abortSignal: AbortSignal.timeout(120_000),
  }
)) {
  if (chunk.version !== version) {
    answer = '';
    version = chunk.version;
  }
  answer += chunk.delta.answer ?? '';
}
if (!steered || !answer.includes('Small launch'))
  throw new Error('Steering was not reflected in the final answer');
console.log(answer);

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

const llm = ai({
  name: 'openai',
  apiKey: process.env.OPENAI_API_KEY ?? process.env.OPENAI_APIKEY,
  config: { model: AxAIOpenAIModel.GPT6Astra, maxTokens: 2000 },
});
const control = runControl();
let overlap = false;
let lookupPending = false;
let calculationOverlapped = false;
let lookupCalls = 0;
let calculationCalls = 0;
control.onEvent((event) => {
  if (event.type === 'model.output' && event.pendingCallIds?.length) {
    overlap = true;
  }
});
const calculate = fn('calculateEstimate')
  .description('Compute an independent estimate while delivery lookup runs')
  .execution('background')
  .handler(() => {
    calculationCalls++;
    calculationOverlapped ||= lookupPending;
    return 7 * 8;
  })
  .build();
const lookup = fn('lookupDelivery')
  .description(
    'Get the current delivery code. Call once; while it runs, reason about a concise customer update.'
  )
  .execution('background')
  .handler(async () => {
    lookupCalls++;
    lookupPending = true;
    console.log('Delivery lookup started');
    await new Promise((resolve) => setTimeout(resolve, 8000));
    lookupPending = false;
    console.log('Delivery lookup completed');
    return { deliveryCode: 'AX-742', status: 'arrives tomorrow' };
  })
  .build();
const result = await ax('question -> answer', {
  functions: [lookup, calculate],
}).forward(
  llm,
  {
    question:
      'Start lookupDelivery first, then call calculateEstimate while lookupDelivery is still running. Include the estimate and delivery code/status in the final answer.',
  },
  {
    control,
    thinkingTokenBudget: 'low',
    serviceTier: 'standard',
    abortSignal: AbortSignal.timeout(120_000),
  }
);
if (!['AX-742', '56', 'tomorrow'].every((text) => result.answer.includes(text)))
  throw new Error('Final answer omitted the tool result');
if (lookupCalls !== 1 || calculationCalls !== 1)
  throw new Error('A tool was not executed exactly once');
if (!overlap || !calculationOverlapped)
  throw new Error(
    'No independent model output observed while a tool was pending'
  );
console.log('Observed independent model work while a tool was pending');
console.log(result.answer);

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

const llm = ai({
  name: 'openai',
  apiKey: process.env.OPENAI_API_KEY ?? process.env.OPENAI_APIKEY,
  config: { model: AxAIOpenAIModel.GPT56Luna },
});
let calls = 0;
const lookup = fn('lookupCode')
  .description('Get the confirmation code')
  .execution('background')
  .handler(() => {
    calls++;
    return 'AX-742';
  })
  .build();
const result = await ax('question -> answer', { functions: [lookup] }).forward(
  llm,
  { question: 'Use lookupCode and report the confirmation code.' },
  {
    thinkingTokenBudget: 'none',
    serviceTier: 'standard',
    abortSignal: AbortSignal.timeout(120_000),
  }
);
if (calls !== 1 || !result.answer.includes('AX-742'))
  throw new Error(
    'Fallback did not execute and incorporate the tool exactly once'
  );
console.log(result.answer);

import { AxAIOpenAIModel, ai, ax } 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 program = ax(
  'context:string, question:string -> answer:string, citations:string[]'
);
const result = await program.forward(llm, {
  context:
    'Ax uses signatures for typed IO, ai() for providers, ax() for generation, agent() for runtime loops, flow() for orchestration, and optimize() for GEPA tuning.',
  question: 'How should a new developer think about Ax?',
});

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

import type { AxRateLimiterFunction, AxRuntimeHooks } from '@ax-llm/ax';
import { AxAIOpenAIModel, agent, ai, ax, axGlobals, flow } from '@ax-llm/ax';
import type { Meter, Span, Tracer } from '@opentelemetry/api';

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 limiter =
  (label: string): AxRateLimiterFunction =>
  async (next, info) => {
    console.log(
      `[limit:${label}] ${info.operation} ${info.provider}/${info.model} stream=${info.streaming}`
    );
    return next();
  };

const span = (name: string): Span =>
  ({
    addEvent: (event: string) => console.log(`[span:event] ${name} ${event}`),
    addLink: () => undefined,
    addLinks: () => undefined,
    end: () => console.log(`[span:end] ${name}`),
    isRecording: () => true,
    recordException: (error: unknown) =>
      console.log(`[span:error] ${name} ${String(error)}`),
    setAttribute: () => undefined,
    setAttributes: () => undefined,
    setStatus: () => undefined,
    spanContext: () => ({
      traceId: '0'.repeat(32),
      spanId: '0'.repeat(16),
      traceFlags: 0,
    }),
    updateName: () => undefined,
  }) as unknown as Span;

const tracer = {
  startSpan: (name: string) => {
    console.log(`[span:start] ${name}`);
    return span(name);
  },
  startActiveSpan: async (name: string, ...args: unknown[]) => {
    console.log(`[span:start] ${name}`);
    const callback = args.at(-1) as (active: Span) => Promise<unknown>;
    return callback(span(name));
  },
} as Tracer;

const instrument = (name: string) => ({
  add: (value: number) => console.log(`[metric] ${name} += ${value}`),
  record: (value: number) => console.log(`[metric] ${name} = ${value}`),
});
const meter = {
  createCounter: (name: string) => instrument(name),
  createGauge: (name: string) => instrument(name),
  createHistogram: (name: string) => instrument(name),
  createObservableCounter: () => ({}),
  createObservableGauge: () => ({}),
  createObservableUpDownCounter: () => ({}),
  createUpDownCounter: (name: string) => instrument(name),
  createBatchObservableCallback: () => ({ dispose() {} }),
  removeBatchObservableCallback: () => undefined,
} as unknown as Meter;

const llm = ai({
  name: 'openai',
  apiKey,
  config: { model: AxAIOpenAIModel.GPT54Mini, temperature: 0 },
});
const overrideHooks: AxRuntimeHooks = {
  rateLimiter: limiter('forward'),
  tracer,
  meter,
};

axGlobals.rateLimiter = limiter('global');
axGlobals.tracer = tracer;
axGlobals.meter = meter;
try {
  const direct = ax('topic:string -> summary:string');
  console.log(
    await direct.forward(llm, { topic: 'portable Ax runtime hooks' })
  );

  const helper = agent('question:string -> answer:string', {});
  console.log(
    await helper.forward(
      llm,
      { question: 'What does a rate limiter wrap?' },
      overrideHooks
    )
  );

  const workflow = flow(`flowchart TD
    %%ax outline: topic:string -> outline:string
    %%ax polish: outline:string -> answer:string
    outline --> polish`);
  console.log(
    await workflow.forward(llm, { topic: 'Ax runtime hooks' }, overrideHooks)
  );
} finally {
  axGlobals.rateLimiter = undefined;
  axGlobals.tracer = undefined;
  axGlobals.meter = undefined;
}

import {
  AxAIAnthropicModel,
  AxAIOpenAIModel,
  AxBalancer,
  AxInMemoryBalancerStatsStore,
  ai,
  ax,
} from '@ax-llm/ax';

const openaiKey = process.env.OPENAI_API_KEY ?? process.env.OPENAI_APIKEY;
const anthropicKey =
  process.env.ANTHROPIC_API_KEY ?? process.env.ANTHROPIC_APIKEY;
if (!openaiKey || !anthropicKey) {
  throw new Error(
    'Set OPENAI_API_KEY (or OPENAI_APIKEY) and ANTHROPIC_API_KEY (or ANTHROPIC_APIKEY).'
  );
}

const openai = ai({
  name: 'openai',
  apiKey: openaiKey,
  models: [
    {
      key: 'fast',
      model: AxAIOpenAIModel.GPT54Mini,
      description: 'Fast general-purpose model',
    },
  ],
});

const anthropic = ai({
  name: 'anthropic',
  apiKey: anthropicKey,
  models: [
    {
      key: 'fast',
      model: AxAIAnthropicModel.Claude45Haiku,
      description: 'Fast general-purpose model',
    },
  ],
});

// Reuse this store across balancers in one process. For multiple processes,
// provide an AxBalancerStatsStore backed by Redis or your application database.
const statsStore = new AxInMemoryBalancerStatsStore();
const routeKeys = new Map<string, string>([
  [openai.getId(), 'openai-primary'],
  [anthropic.getId(), 'anthropic-primary'],
]);

const llm = AxBalancer.create([openai, anthropic] as const, {
  strategy: {
    type: 'adaptive',
    deadlineMs: 6_000,
    badOutcomeCost: 0.02,
    expectedTokens: { promptTokens: 1_200, completionTokens: 300 },
    namespace: 'support-summary-v1',
    routeKey: (service) => {
      const key = routeKeys.get(service.getId());
      if (!key) throw new Error('Missing stable route key.');
      return key;
    },
    slice: ({ options }) =>
      options?.customLabels?.workflow ?? 'default-workflow',
    statsStore,
    // Analytics only: statsStore remains the authoritative decision state.
    onRoutingEvent: (event) => {
      if (event.type === 'selected' || event.type === 'fallback') {
        console.log('route:', event);
      }
    },
  },
});

const summarize = ax('supportTicket:string -> summary:string, urgency:string');
const result = await summarize.forward(
  llm,
  {
    supportTicket:
      'Our checkout started timing out after the latest deployment.',
  },
  {
    model: 'fast',
    customLabels: { workflow: 'support-summary' },
  }
);

console.log(result);

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

const control = runControl();
const completed: string[] = [];
const counts = { slow: 0, fast: 0 };
let slowPending = false;
let overlap = false;
control.onEvent((event) => {
  if (event.type === 'applied') console.log(`Update applied: ${event.timing}`);
});
const slow = fn('lookupSlowCode')
  .description('Get the slow code. Call exactly once, before lookupFastCode.')
  .execution('background')
  .handler(async () => {
    counts.slow++;
    slowPending = true;
    await new Promise((resolve) => setTimeout(resolve, 12_000));
    slowPending = false;
    completed.push('slow');
    return 'SLOW-42';
  })
  .build();
const fast = fn('lookupFastCode')
  .description(
    'Get the independent fast code. Call exactly once while lookupSlowCode runs.'
  )
  .execution('background')
  .handler(async () => {
    counts.fast++;
    overlap ||= slowPending;
    control.steer('Report both exact codes and end the answer with Confirmed.');
    control.setThinkingTokenBudget('medium');
    await new Promise((resolve) => setTimeout(resolve, 100));
    control.steer(
      'Also include the word VERIFIED, while retaining both exact codes and ending with Confirmed.'
    );
    completed.push('fast');
    return 'FAST-17';
  })
  .build();
const llm = ai({
  name: 'openai',
  apiKey: process.env.OPENAI_API_KEY ?? process.env.OPENAI_APIKEY,
  config: { model: AxAIOpenAIModel.GPT6Astra, maxTokens: 2500 },
  options: { webSocket: WebSocket },
});
const result = await ax('question -> answer', {
  functions: [slow, fast],
}).forward(
  llm,
  {
    question:
      'Start lookupSlowCode, then immediately call lookupFastCode while the slow lookup runs. Use both results in a short final answer.',
  },
  {
    control,
    thinkingTokenBudget: 'low',
    serviceTier: 'standard',
    abortSignal: AbortSignal.timeout(120_000),
  }
);
if (!overlap || completed.join(',') !== 'fast,slow')
  throw new Error('Expected independent, out-of-order tool completion');
if (counts.fast !== 1 || counts.slow !== 1)
  throw new Error('A tool was executed more than once');
if (
  !['FAST-17', 'SLOW-42', 'Confirmed', 'VERIFIED'].every((text) =>
    result.answer.includes(text)
  )
)
  throw new Error('Final answer omitted a result or the steering instruction');
console.log(`Completed tools: ${completed.join(', ')}`);
console.log(result.answer);