Ax Llm

---
name: ax-llm
description: This skill helps with using the @ax-llm/ax TypeScript library for building LLM applications. Use when the user asks about ax(), ai(), f(), s(), agent(), flow(), AxGen, AxAgent, AxFlow, signatures, streaming, or mentions @ax-llm/ax.
version: "25.1.0"
---

# Ax Library (@ax-llm/ax) Quick Reference

Ax is a TypeScript library for building LLM-powered applications with type-safe signatures, streaming support, and multi-provider compatibility.

> **Detailed skills available:** ax-ai (providers, routing, adaptive balancing), ax-typesafe (Typesafe/Jev decisions, criteria, native scoring), ax-signature (signatures/types), ax-gen (generators), ax-agent (core agents/tools), ax-agent-rlm (agent runtime/RLM/delegation), ax-agent-observability (callbacks/logs/usage), ax-agent-memory-skills (recall and dynamic skill loading), ax-agent-optimize (agent tuning/eval), ax-flow (workflows), ax-gepa (top-level `optimize(...)`, BootstrapFewShot -> GEPA, Pareto optimization).

## Imports & Factories

```typescript
// Prefer factory functions: ax(), ai(), agent(), flow(); avoid class constructors.
import { ax, ai, f, s, fn, agent, flow, AxMemory, AxMCPClient } from '@ax-llm/ax';
import { z } from 'zod'; // optional — any Standard Schema v1 library works

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

// Generator (from string signature)
const gen = ax('question:string -> answer:string');

// Generator (from fluent signature)
const fluentGen = ax(
  f()
    .input('question', f.string('User question'))
    .output('answer', f.string('AI response'))
    .build()
);

// Generator (from zod — Standard Schema v1, also works with valibot/arktype)
const zodGen = ax(
  f()
    .input(z.object({ question: z.string().describe('User question') }))
    .output(z.object({ answer: z.string().describe('AI response') }))
    .build()
);

// Reusable signature
const sig = s('question:string, context:string[] -> answer:string');

// Agent
const myAgent = agent('userInput:string -> response:string', {
  agentIdentity: {
    name: 'helper',
    description: 'A helpful assistant',
  },
});

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

// Function tool — native fluent
const tool = fn('search')
  .description('Search the web')
  .arg('query', f.string('Search query'))
  .returns(f.string('Search results'))
  .handler(({ query }) => searchWeb(query))
  .build();

// Function tool — zod schema (Standard Schema v1: also works with valibot, arktype)
const zodTool = fn('calculateTax')
  .description('Calculate tax for an amount')
  .arg(z.object({
    amount: z.number().positive().describe('Pre-tax amount in USD'),
    region: z.enum(['US', 'EU', 'UK']).describe('Tax region'),
  }))
  .returns(z.object({ tax: z.number(), total: z.number() }))
  .handler(async ({ amount }) => ({ tax: amount * 0.1, total: amount * 1.1 }))
  .build();
```

## Running

```typescript
// Forward (blocking)
const result = await gen.forward(llm, { question: 'What is 2+2?' });

// Streaming
for await (const chunk of gen.streamingForward(llm, { question: 'Tell a story' })) {
  if (chunk.delta.answer) process.stdout.write(chunk.delta.answer);
}
```

## Forward Options Quick Reference

Provider capabilities determine which outputs and options apply. Typesafe/Jev
supports required boolean/class signatures; use its separate native client for
Score and full probabilities, and a generative provider for prose or tools. See
the [ax-typesafe skill](https://github.com/ax-llm/ax/blob/main/src/ax/skills/ax-typesafe.md)
before applying the general generation options below to Jev.

| Goal | Option | Example |
|------|--------|---------|
| Model override | `model` | `{ model: 'gpt-5.4-mini' }` |
| Temperature | `modelConfig.temperature` | `{ modelConfig: { temperature: 0.8 } }` |
| Max tokens | `modelConfig.maxTokens` | `{ modelConfig: { maxTokens: 500 } }` |
| Retry on failure | `maxRetries` | `{ maxRetries: 3 }` |
| Max agent steps | `maxSteps` | `{ maxSteps: 10 }` |
| Fail fast | `fastFail` | `{ fastFail: true }` |
| Thinking budget | `thinkingTokenBudget` | `{ thinkingTokenBudget: 'medium' }` |
| Show thoughts | `showThoughts` | `{ showThoughts: true }` |
| Context caching | `contextCache` | `{ contextCache: { cacheBreakpoint: 'after-examples' } }` |
| Multi-sampling | `sampleCount` | `{ sampleCount: 5 }` |
| Debug logging | `debug` | `{ debug: true }` |
| Abort signal | `abortSignal` | `{ abortSignal: controller.signal }` |
| Memory | `mem` | `{ mem: new AxMemory() }` |
| Stop function | `stopFunction` | `{ stopFunction: 'finalAnswer' }` |
| Function mode | `functionCallMode` | `{ functionCallMode: 'auto' }` |

Global runtime defaults can be set with `axGlobals` and are read live by future AI, AxGen, and AxFlow calls:

```typescript
import { axGlobals, axCreateDefaultColorLogger } from '@ax-llm/ax';
import { metrics, trace } from '@opentelemetry/api';

axGlobals.rateLimiter = async (next, info) => next();
axGlobals.tracer = trace.getTracer('my-app');
axGlobals.meter = metrics.getMeter('my-app');
axGlobals.debug = true;
axGlobals.logger = axCreateDefaultColorLogger();
```

Runtime hooks resolve as: forward/direct-call hooks, enclosing program defaults, child-program defaults, AI-service hooks, then globals snapshotted at operation start. They are native run-scoped values and never enter AxIR JSON state, cache keys, exported state, traces, or optimizer artifacts. Agent and flow forwards carry them through every internal generator and model call without mutating children or leaking across concurrent runs. Limiter failures propagate; tracer, meter, and usage-observer failures are fail-open. `customLabels` merge by precedence, and `abortSignal` values are combined so either global or local cancellation works.

## Memory and Context

```typescript
import { AxMemory } from '@ax-llm/ax';

const memory = new AxMemory();

// Multi-turn conversation
await gen.forward(llm, { userMessage: 'My name is Alice' }, { mem: memory });
const r = await gen.forward(llm, { userMessage: 'What is my name?' }, { mem: memory });
```

## Few-Shot Examples

```typescript
const classifier = ax('reviewText:string -> sentiment:class "positive, negative, neutral"');

classifier.setExamples([
  { reviewText: 'I love this!', sentiment: 'positive' },
  { reviewText: 'Terrible.', sentiment: 'negative' },
  { reviewText: 'It works.', sentiment: 'neutral' },
]);
```

## Common Patterns

### Classification

```typescript
const classifier = ax(
  f()
    .input('text', f.string())
    .output('category', f.class(['spam', 'ham', 'uncertain']))
    .output('confidence', f.number().min(0).max(1))
    .build()
);
```

### Extraction

```typescript
const extractor = ax(
  f()
    .input('text', f.string())
    .output('entities', f.object({
      people: f.string().array(),
      organizations: f.string().array(),
      locations: f.string().array()
    }))
    .build()
);
```

### Multi-modal (Images)

```typescript
const analyzer = ax(
  f()
    .input('image', f.image('Image to analyze'))
    .input('question', f.string('Question').optional())
    .output('description', f.string())
    .output('objects', f.string().array())
    .build()
);

const result = await analyzer.forward(llm, {
  image: { mimeType: 'image/jpeg', data: base64Data },
  question: 'What objects are in this image?'
});
```

### Chaining Generators

```typescript
const researcher = ax('topic:string -> research:string, keyFacts:string[]');
const writer = ax('research:string, keyFacts:string[] -> article:string');

const research = await researcher.forward(llm, { topic: 'AGI' });
const draft = await writer.forward(llm, { research: research.research, keyFacts: research.keyFacts });
```

## Error Handling

```typescript
import { AxGenerateError, AxAIServiceError, AxAIServiceAbortedError } from '@ax-llm/ax';

try {
  const result = await gen.forward(llm, { input: 'test' });
} catch (error) {
  if (error instanceof AxGenerateError) {
    console.error('Generation failed:', error.details.model, error.details.signature);
  } else if (error instanceof AxAIServiceAbortedError) {
    console.log('Request was aborted');
  } else if (error instanceof AxAIServiceError) {
    console.error('AI service error:', error.message);
  }
}
```

## Debugging

```typescript
import { axCreateDefaultColorLogger, axGlobals } from '@ax-llm/ax';

const result = await gen.forward(llm, { input: 'test' }, {
  debug: true,
  logger: axCreateDefaultColorLogger(),
  // OpenTelemetry
  tracer: openTelemetryTracer,
  meter: openTelemetryMeter,
});

// Or set live app-wide defaults for future calls:
axGlobals.tracer = openTelemetryTracer;
axGlobals.meter = openTelemetryMeter;
```

## MCP Integration

Use the `ax-mcp` skill for the complete native client, transport,
authentication, catalog, task, subscription, event, and replay workflow.

```typescript
import { AxMCPClient, agent } from '@ax-llm/ax';
import { AxMCPStdioTransport } from '@ax-llm/ax-tools';

// Stdio transport (local MCP server)
const transport = new AxMCPStdioTransport({
  command: 'npx',
  args: ['-y', '@modelcontextprotocol/server-memory'],
});

const mcpClient = new AxMCPClient(transport, { namespace: 'memory' });

// Native MCP context is initialized once and inherited by all agent stages.
const myAgent = agent('userMessage:string -> response:string', {
  mcp: mcpClient,
  functionDiscovery: true,
  contextFields: [],
});

const result = await myAgent.forward(llm, { userMessage: 'Remember this.' });
await mcpClient.close(); // caller-owned clients remain caller-owned
```

### HTTP Transport (Remote MCP)

```typescript
import { AxMCPStreamableHTTPTransport } from '@ax-llm/ax';

const transport = new AxMCPStreamableHTTPTransport('https://remote.example/mcp', {
  headers: { 'x-pd-project-id': projectId },
  authorization: `Bearer ${accessToken}`,
});
```

### Native MCP and UCP behavior

- Pass `mcp` and `ucp` to AxGen, streaming AxGen, chat, AxAgent, AxFlow, optimization, or evaluation options.
- Use `mcpContext` to inject attributed prompts/resources before the first model call.
- Use `mcpInheritance: 'all' | 'none' | string[]` to restrict child programs.
- Tool calls retain raw MCP content, metadata, errors, tasks, and protocol provenance in memory.
- AxAgent exposes native modules as `mcp.<namespace>` and `ucp.<namespace>`.
- `inspectCatalog()` discovers tool/prompt names, concrete resources, and URI
  templates from only an endpoint. Resource event sources default to no
  subscriptions and require an explicit all/URI/selector policy.
- `toFunction()` remains a compatibility adapter only; native Ax execution never uses it.
- Live optimization is rejected by default. Use recording/replay or explicitly opt into live MCP evaluation.

```typescript
const catalog = await mcpClient.inspectCatalog();
const tools = catalog.tools;
const prompts = await mcpClient.listPrompts();
const resource = await mcpClient.readResource('docs://guide');
const tasks = await mcpClient.listTasks();
```

### Function Overrides

```typescript
const mcpClient = new AxMCPClient(transport, {
  functionOverrides: [
    { name: 'search_documents', updates: { name: 'findDocs', description: 'Search docs' } }
  ]
});
```

## Type Reference

All programs use `forward(ai, values, options?)` and
`streamingForward(ai, values, options?)`. For an AI service type `T`, use the
exported option types below instead of constructor options:

| Program | Forward options | Streaming options |
| --- | --- | --- |
| `AxGen` | `AxProgramForwardOptionsWithModels<T>` | `AxProgramStreamingForwardOptionsWithModels<T>` |
| `AxAgent` | `AxAgentForwardOptions<T>` | `AxAgentStreamingForwardOptions<T>` |
| `AxFlow` | `AxFlowForwardOptions<T>` | `AxProgramStreamingForwardOptionsWithModels<T>` |

`forward()` returns `Promise<OUT>`. `streamingForward()` returns
`AxGenStreamingOut<OUT>` with `{ version, index, delta }` chunks. Agent input
values can also include a per-run `memories` array. `AxAgentOptions<IN>` is for
constructing an agent, not for its forward call.

- `AxGen`: `setExamples(...)`, `addAssert(...)`, `addStreamingAssert(...)`,
  `addFieldProcessor(...)`, `addStreamingFieldProcessor(...)`, and `stop()`.
  See `ax-gen` for their callback contracts.
- `AxAgent`: `getFunction()` exposes a child-agent tool and requires
  `agentIdentity` at construction. See `ax-agent`.
- `AxFlow`: define nodes with `.node(...)`, execute them with `.execute(...)`,
  and select outputs with `.returns(...)`. See `ax-flow`.

## Event-Driven Programs

Use `eventRuntime()` when notifications, webhooks, timers, or remote tasks
should wake or resume an Ax program. Sources publish into an inbox; explicit
routes choose `observe`, `invalidate`, `wake`, or `resume`. Event payloads are
never inserted as user messages automatically. See `ax-event-runtime.md`.

## Examples

Fetch these for full working code:

- [Standard Schema (zod)](https://raw.githubusercontent.com/ax-llm/ax/refs/heads/main/src/examples/standard-schema.ts) — zod with f() and fn()
- [Chat Log](https://raw.githubusercontent.com/ax-llm/ax/refs/heads/main/src/examples/chat-log.ts) — inspect conversation history and multi-step tool calls
- [Marketing](https://raw.githubusercontent.com/ax-llm/ax/refs/heads/main/src/examples/marketing.ts) — product use case
- [MCP Integration](https://raw.githubusercontent.com/ax-llm/ax/refs/heads/main/src/examples/mcp-client-memory.ts) — MCP integration