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