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 { 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, 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 { 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.Gemini37Flash;
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 } 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);