Advanced Start

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, AxJSRuntime, agent, ai } 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,
  },
});

// The handbook can be arbitrarily large. Listing it in `contextFields` keeps it
// in the agent's runtime so it never inflates the model prompt — the agent reads
// it through code, not through tokens. That is the whole point of an Ax agent
// over a plain `ax()` call: the source material stays out of the context window.
const handbook = `
# Acme Cloud — Support Handbook

## Billing
- Invoices are issued on the 1st of each month and are due net-15.
- Plan downgrades take effect at the END of the current billing cycle, not immediately.
- Refunds are issued to the original payment method within 5 business days.

## Access
- Seats can be added by any workspace Owner under Settings -> Members.
- SSO (SAML) is available on Enterprise; SCIM provisioning is Owner-only.

## Incidents
- Status and uptime are published at status.acme.example.
- Sev-1 incidents page the on-call within 5 minutes; updates post every 30 minutes.

## Data
- Exports are available in CSV and JSON from Settings -> Data.
- Deleted workspaces are recoverable for 30 days, then permanently purged.
`.trim();

const assistant = agent(
  'question:string, handbook:string -> answer:string, citations:string[] "Handbook sections the answer relies on"',
  {
    runtime: new AxJSRuntime(),
    // Keep the handbook in the runtime, out of the prompt.
    contextFields: ['handbook'],
    maxTurns: 6,
  }
);

const result = await assistant.forward(llm, {
  question:
    'A customer downgraded their plan today. When does it take effect, and can they get a refund for the current cycle?',
  handbook,
});

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<{ 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 { writeFileSync } from 'node:fs';
import { AxAIOpenAIModel, ai } 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 speech = await llm.speak({
  text: 'Ax turns LLM prompts into typed programs.',
  voice: 'alloy',
  format: 'mp3',
});

writeFileSync(
  './src/examples/output/public-openai-speech.mp3',
  Buffer.from(speech.data, 'base64')
);
console.log(
  JSON.stringify(
    { format: speech.format, transcript: speech.transcript },
    null,
    2
  )
);

import { AxAIOpenAIModel, ai, ax, optimize } 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(
  'emailText:string -> priority:class "high, normal, low", rationale:string'
);

const baseline = await program.forward(llm, {
  emailText: 'Production checkout is failing for enterprise customers.',
});

const train = [
  {
    emailText: 'URGENT: checkout is down',
    priority: 'high',
    rationale: 'Production checkout outage blocks customers.',
  },
  {
    emailText: 'Weekly newsletter',
    priority: 'low',
    rationale: 'Informational update with no action needed.',
  },
  {
    emailText: 'Reminder to submit timesheets',
    priority: 'normal',
    rationale: 'Routine request with a clear deadline.',
  },
];

const metric = ({ prediction, example }: { prediction: any; example: any }) =>
  prediction.priority === example.priority ? 1 : 0;

const result = await optimize(program, train, metric, {
  studentAI: llm,
  teacherAI: llm,
  numTrials: 1,
  maxMetricCalls: 4,
});

if (!result.optimizedProgram) {
  throw new Error('Optimizer did not return an optimized program.');
}

program.applyOptimization(result.optimizedProgram);
const after = await program.forward(llm, {
  emailText: 'Production checkout is failing for enterprise customers.',
});

console.log(
  JSON.stringify({ baseline, after, bestScore: result.bestScore }, null, 2)
);

import {
  AxAIOpenAIModel,
  AxMCPClient,
  AxMCPStreamableHTTPTransport,
  ai,
  ax,
} from '@ax-llm/ax';
import { AxMCPEventDemoServer } from '../../mcp-event-demo-server.js';

const apiKey = process.env.OPENAI_API_KEY ?? process.env.OPENAI_APIKEY;
if (!apiKey) throw new Error('Set OPENAI_API_KEY or OPENAI_APIKEY.');

const server = new AxMCPEventDemoServer();
const endpoint = await server.start();
const mcp = new AxMCPClient(
  new AxMCPStreamableHTTPTransport(endpoint, {
    ssrfProtection: { allowHTTP: true, allowLoopback: true },
  }),
  { namespace: 'inventory' }
);
const llm = ai({
  name: 'openai',
  apiKey,
  config: { model: AxAIOpenAIModel.GPT54Mini, temperature: 0 },
});
const program = ax(
  'taskRequest:string -> answer:string "Use the inventory MCP tool and report its task id."',
  { mcp }
);

try {
  const catalog = await mcp.inspectCatalog();
  console.log({
    tools: catalog.tools.map(({ name }) => name),
    resources: catalog.resources.map(({ name, uri }) => ({ name, uri })),
    resourceTemplates: catalog.resourceTemplates.map(
      ({ name, uriTemplate }) => ({ name, uriTemplate })
    ),
  });
  console.log(
    await program.forward(llm, { taskRequest: 'Reindex inventory.' })
  );
} finally {
  await mcp.close();
  await server.close();
}

import {
  AxAIOpenAIModel,
  AxJSRuntime,
  AxMCPClient,
  AxMCPEventSource,
  AxMCPStreamableHTTPTransport,
  agent,
  ai,
  eventPath,
  eventRoute,
  eventRuntime,
  eventTarget,
} from '@ax-llm/ax';
import {
  AxMCPEventDemoServer,
  waitForDemoSignal,
} from '../../mcp-event-demo-server.js';

const apiKey = process.env.OPENAI_API_KEY ?? process.env.OPENAI_APIKEY;
if (!apiKey) throw new Error('Set OPENAI_API_KEY or OPENAI_APIKEY.');

const server = new AxMCPEventDemoServer();
const endpoint = await server.start();
const mcp = new AxMCPClient(
  new AxMCPStreamableHTTPTransport(endpoint, {
    ssrfProtection: { allowHTTP: true, allowLoopback: true },
  }),
  { namespace: 'inventory' }
);
const llm = ai({
  name: 'openai',
  apiKey,
  config: { model: AxAIOpenAIModel.GPT54Mini },
});
const program = agent('uri:string -> summary:string', {
  runtime: new AxJSRuntime(),
});
let completeWake!: () => void;
const wakeCompleted = new Promise<void>((resolve) => {
  completeWake = resolve;
});
const target = eventTarget('inventory-agent')
  .program(program)
  .ai(llm)
  .input((input) => input.field('uri', eventPath.data('uri')))
  .forwardOptions({ mcp })
  .sink({
    id: 'console',
    write: (output) => {
      console.log(output);
      completeWake();
    },
  })
  .build();
const runtime = eventRuntime({
  allowVolatile: true,
  sources: [
    new AxMCPEventSource({
      client: mcp,
      resourceSubscriptions: 'all',
      identity: { tenantId: 'demo' },
      trust: 'authenticated',
    }),
  ],
  routes: [
    eventRoute('resource-wake')
      .types('mcp.resource.updated')
      .authenticated()
      .instanceKey(eventPath.subject())
      .wake(target)
      .build(),
  ],
});

try {
  await runtime.start();
  await server.waitForListeningConnection();
  await server.waitForSubscription('demo://inventory');
  server.updateResource();
  await waitForDemoSignal(wakeCompleted, 'the MCP resource wake');
  await runtime.waitForIdle();
} finally {
  await runtime.close({ drain: false });
  await mcp.close();
  await server.close();
}

export * from '../../mcp-task-resume-flow.js';