Ax Gen

AxGen Codegen Rules (@ax-llm/ax)

Use this skill to generate AxGen code. Prefer short, modern, copyable patterns. Do not write tutorial prose unless the user explicitly asks for explanation.

Use the ax-mcp skill when AxGen attaches native MCP clients or consumes MCP prompts, resources, tools, tasks, subscriptions, authentication, or events.

Use These Defaults

• Use ax(...) factory, not new AxGen(...).
• Always pass an AI instance from ai(...) as the first argument to forward().
• Streaming uses streamingForward(), not forward() with a stream option.
• Use schema validation for field shape and constraints.
• Use addAssert(...) for whole-output hard invariants with correction retries.
• Use addStreamingAssert(...) for partial streaming hard invariants with fail-fast per-attempt correction retries.
• Use bestOfN(...) / refine(...) for reward-scored complete outputs.
• Step hook mutations are applied at the next step boundary (pending pattern).
• stopFunction accepts a string or string[] for multiple stop functions.
• Multi-step continues until: all outputs filled, stop function called, or maxSteps reached.

Canonical Pattern

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

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

// Inline signature
const gen = ax('input:string -> output:string, reasoning:string');

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

// With options
const gen3 = ax('input -> output', {
  description: 'A helpful assistant',
  maxRetries: 3,
  maxSteps: 10,
  temperature: 0.7,
});

const result = await gen.forward(llm, { input: 'Hello world' });
console.log(result.output);

Signatures from zod / valibot / arktype

ax() accepts any signature built with f(), and f().input() / .output() accept Standard Schema v1 validators directly — per-field or a whole z.object({...}):

import { z } from 'zod';
import { ax, f } from '@ax-llm/ax';

const gen = ax(
  f()
    .input(z.object({
      productName: z.string(),
      buyerProfile: z.string(),
    }))
    .output(z.object({
      headline: z.string(),
      recommendation: z.enum(['buy', 'wait', 'skip']),
    }))
    .build()
);

Constraints (.min(), .email(), .regex()) and custom logic (.refine(), .transform(), .superRefine()) execute in the normal validation/retry pipeline — at parse time on complete field values, including at field boundaries during streaming. For cache/internal hints pass companion options: .input('ctx', z.string(), { cache: true }) or .output('reasoning', z.string(), { internal: true }).

Define tool functions with zod the same way — fn().arg() / .returns() accept per-argument or whole-object schemas and infer the handler's argument type:

import { z } from 'zod';
import { ax, fn } from '@ax-llm/ax';

const lookupProduct = fn('lookupProduct')
  .description('Look up a product by name')
  .arg(z.object({
    productName: z.string().min(1),
    includeSpecs: z.boolean().optional(),
  }))
  .returns(z.object({
    price: z.number(),
    inStock: z.boolean(),
    rating: z.number().min(1).max(5),
  }))
  .handler(async ({ productName, includeSpecs }) => ({
    price: 79.99,
    inStock: true,
    rating: 4.3,
  }))
  .build();

const result = await gen.forward(llm, { ... }, { functions: [lookupProduct] });

Running AxGen

forward()

const result = await gen.forward(llm, { input: '...' });

// With options
const result = await gen.forward(llm, { input: '...' }, {
  maxRetries: 5,
  model: 'gpt-5.4-mini',
  modelConfig: { temperature: 0.9, maxTokens: 1000 },
  debug: true,
});

Live Global Defaults

AxGen respects axGlobals for app-wide runtime defaults:

import { axGlobals } from '@ax-llm/ax';
import { trace } from '@opentelemetry/api';

const responseCache = new Map<string, any>();

axGlobals.rateLimiter = async (next, info) => next();
axGlobals.tracer = trace.getTracer('my-app');
axGlobals.debug = true;
axGlobals.cachingFunction = async (key, value?) => {
  if (value !== undefined) {
    responseCache.set(key, value);
    return;
  }
  return responseCache.get(key);
};

Rules:

• Runtime-hook precedence is: forward options, then generator options, then AI service options, then the globals snapshotted at run start.
• A forward-scoped rateLimiter, tracer, or meter is carried to every retry and provider call without being serialized or mutating the generator. Concurrent forwards remain isolated.
• Limiter failures propagate. Tracer and meter failures are ignored, and telemetry contains metadata rather than prompts, outputs, or tool payloads.
• abortSignal from axGlobals is merged with local forward signals.
• customLabels merge from globals to AI service to forward options.
• cachingFunction and functionResultFormatter also fall back to current axGlobals when local options do not provide them.

streamingForward()

const stream = gen.streamingForward(llm, { input: 'Write a long story' });
for await (const chunk of stream) {
  if (chunk.delta.output) process.stdout.write(chunk.delta.output);
}

Stopping And Cancellation

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

const timer = setTimeout(() => gen.stop(), 3_000);

try {
  const result = await gen.forward(llm, { topic: 'Long document' }, {
    abortSignal: AbortSignal.timeout(10_000),
  });
} catch (err) {
  if (err instanceof AxAIServiceAbortedError) console.log('Aborted');
}

Rules:

• gen.stop() gracefully stops multi-step execution at the next step boundary.
• abortSignal cancels the underlying AI service call immediately.
• Catch AxAIServiceAbortedError when using either mechanism.

Validation, Selection, And Guards

import { ax, bestOfN, f } from '@ax-llm/ax';
import { z } from 'zod';

// Schema validation: output shape and field validity.
const gen = ax(
  f()
    .input('topic', z.string().min(1))
    .output('summary', z.string().min(50))
    .build()
);

// bestOfN: choose the best complete candidate.
const selected = bestOfN(gen, {
  n: 4,
  rewardFn: ({ prediction }) => prediction.summary.length,
});

// Whole-output assertion: retries with correction feedback.
gen.addAssert(
  (output) => output.summary.includes(topic) || 'Summary must mention the topic.'
);

// Streaming assertion: fail fast on unsafe partial output.
gen.addStreamingAssert(
  'summary',
  (text) => !text.includes('forbidden'),
  'Output contains forbidden text'
);

Rules:

• Schema validation retries with parser/constraint feedback.
• addAssert(...) checks the complete parsed output after validation/processors and retries with correction feedback on failure.
• bestOfN(...) scores complete candidates and returns the highest reward or first threshold hit.
• refine(...) runs rounds and can feed reward-derived advice into instruction components between rounds.
• addStreamingAssert(...) targets a string/code output field and receives partial text so far.
• Streaming assertions abort the current stream attempt by throwing AxStreamingAssertionError, then feed correction feedback into AxGen retries.

Field Processors

// Post-processing after generation
gen.addFieldProcessor('summary', (value, context) => value.toUpperCase());

// Streaming field processor (called on each chunk)
gen.addStreamingFieldProcessor('content', (partialValue, context) => {
  console.log(`Received ${partialValue.length} chars`);
  return partialValue;
});

Rules:

• addFieldProcessor runs once after the field is fully generated.
• addStreamingFieldProcessor runs on each streaming chunk for the target field.
• Both must return the (possibly transformed) value.

Function Calling

const result = await gen.forward(llm, { question: '...' }, {
  functions: tools,
  functionCallMode: 'auto',
  stopFunction: 'finalAnswer',
});

Rules:

• functionCallMode can be 'auto', 'none', or a specific function name to force.
• stopFunction accepts a string or string[] to halt multi-step on specific function calls.
• Multi-step continues until all outputs filled, stop function called, or maxSteps reached.

Caching

Response Caching

const gen = ax('question:string -> answer:string', {
  cachingFunction: async (key, value?) => {
    if (value !== undefined) {
      await cache.set(key, value);
      return;
    }
    return await cache.get(key);
  },
});

Context Caching

const result = await gen.forward(llm, { question: '...' }, {
  contextCache: { cacheBreakpoint: 'after-examples' },
});

Rules:

• cachingFunction acts as a get/set: called with (key) to read, (key, value) to write.
• contextCache enables AI provider-level prompt caching for long context.
• Provider-facing forward options are merged with constructor defaults before the chat call. This includes promptCacheKey, sessionId, and contextCache in TypeScript and every generated language package; per-call values take precedence.

Sampling And Result Picker

const result = await gen.forward(llm, { question: '...' }, {
  sampleCount: 3,
  resultPicker: async (samples) => {
    // Evaluate each sample and return the index of the best one
    return bestIndex;
  },
});

Rules:

• sampleCount generates multiple completions in parallel.
• resultPicker receives all samples and must return the index of the chosen result.

Extended Thinking

const result = await gen.forward(llm, { question: '...' }, {
  thinkingTokenBudget: 'medium',
  showThoughts: true,
});
console.log(result.thought);

Rules:

• thinkingTokenBudget accepts 'none', 'minimal', 'low', 'medium', 'high', or 'highest'. Provider-specific numeric configuration is only for models such as Gemini 2.5 that expose a numeric thinking budget; Gemini 3 uses model-aware thinking levels instead.
• Set showThoughts: true to include the model's reasoning in result.thought.

Structured Outputs

const sig = f()
  .input('text', f.string())
  .output('summary', f.string())
  .output('metadata', f.json().optional())
  .useStructured()
  .build();

Rules:

• .useStructured() asks providers with native support, including OpenAI, Anthropic, and Gemini, for schema-constrained JSON.
• Output names and shapes are part of the prompt contract as well as the provider schema. Ax renders every exact wire key, required/optional status, type, constraints, and nested shape so capability fallback does not erase the contract.
• structuredOutputMode: 'auto' follows the selected profile/model's ordered structuredOutputModes capability list. Exact caller modelInfo overrides win over profile model rules and defaults.
• Without native schema support, one required non-array string or code output can use json_object plus an exact-shape prompt, client-side validation, and bounded correction retries. This optimized path is provider-neutral and does not require provider-visible tools.
• Richer shapes use the first advertised rung. A json_object selection sends no synthetic __axOutput; Ax keeps the exact-shape prompt, strict parsing, and correction retry.
• Ax advertises only __axOutput. It accepts legacy inbound __finalResult calls so stored trajectories remain replayable, and rejects user functions that collide with either reserved name.
• Use structuredOutputMode: 'native' to require native schema enforcement; Ax reports an error instead of silently weakening that requirement.
• Use structuredOutputMode: 'function' to require the function-argument path; Ax reports an error before sending a request when function calling is unavailable.
• Use structuredOutputMode: 'json_object' to require JSON object mode for rich or singleton output; Ax reports an error before transport when the selected profile/model has not verified it.
• Direct json_schema and json_object chat requests validate their corresponding capabilities independently. structuredOutputs remains the compatibility alias for native JSON Schema only.
• Chat-log provenance records the selected path at providerMetadata.ax.structured_output_rung (native, function, or json_object).
• Native structured-output schemas list every object property in required, set additionalProperties: false on objects, and express optional fields as nullable types.
• Flexible json fields and unshaped object fields are sent as JSON-encoded strings for native structured outputs, then parsed back into normal JavaScript values.

Step Hooks

const result = await gen.forward(llm, values, {
  stepHooks: {
    beforeStep: (ctx) => {
      if (ctx.functionsExecuted.has('complexanalysis')) {
        ctx.setModel('smart');
        ctx.setThinkingBudget('high');
      }
    },
    afterStep: (ctx) => {
      console.log(`Usage: ${ctx.usage.totalTokens} tokens`);
    },
  },
});

AxStepContext Read-Only Properties

• stepIndex - current step number
• maxSteps - configured maximum steps
• isFirstStep - whether this is the first step
• functionsExecuted - Set<string> of function names called so far
• lastFunctionCalls - array of the most recent function call results
• usage - token usage statistics
• state - current step state

AxStepContext Mutators

• setModel(model) - change the model for the next step
• setThinkingBudget(budget) - adjust thinking budget
• setTemperature(temp) - adjust temperature
• setMaxTokens(max) - adjust max output tokens
• setOptions(opts) - set arbitrary forward options
• addFunctions(fns) - add functions for the next step
• removeFunctions(names) - remove functions by name
• stop() - stop multi-step execution

Rules:

• All mutations are pending and applied at the next step boundary.
• beforeStep runs before each LLM call; afterStep runs after.
• Use afterFunctionExecution to react to specific function results.

Self-Tuning

// Simple: enable all self-tuning
const result = await gen.forward(llm, values, { selfTuning: true });

// Granular: pick what to tune
const result = await gen.forward(llm, values, {
  selfTuning: {
    model: true,
    thinkingBudget: true,
    functions: [searchWeb, calculate],
  },
});

Rules:

• selfTuning: true enables automatic model and parameter selection.
• Granular config allows tuning specific aspects independently.
• selfTuning.functions provides a pool of functions the tuner may add or remove per step.

Error Handling

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

try {
  const result = await gen.forward(llm, { input: '...' });
} catch (error) {
  if (error instanceof AxGenerateError) {
    console.log(error.details.model, error.details.signature);
  }
}

Rules:

• AxGenerateError includes details with model and signature for debugging.
• AxAIServiceAbortedError is thrown on cancellation via stop() or abortSignal.

Chat Log and Usage

getChatLog()

After any .forward() or streamingForward() call, gen.getChatLog() returns the full normalized chat history — every ai.chat() round-trip, including the system prompt, all messages, and the model response. The log is reset at the start of each .forward() call. Multi-step generators (with function calls) produce one entry per step.

await gen.forward(llm, { question: 'What is 2+2?' });

for (const entry of gen.getChatLog()) {
  console.log('model:', entry.model);
  for (const msg of entry.messages) {
    console.log(`[${msg.role}]`, msg.content);
  }
  console.log('tokens:', entry.modelUsage?.tokens);
}

Message roles: system, user, assistant, tool. Assistant content uses inline XML:

• <think>...</think> — reasoning/thinking tokens
• <tool_call>\n{...}\n</tool_call> — tool invocations

The system message includes a <tools> JSON block when functions are present.

type AxChatLogMessage =
  | { role: 'system'; content: string }
  | { role: 'user'; content: string }
  | { role: 'assistant'; content: string }
  | { role: 'tool'; name: string; content: string };

type AxChatLogEntry = {
  name?: string;
  model: string;
  messages: AxChatLogMessage[];
  modelUsage?: AxProgramUsage;
};

gen.getChatLog(): readonly AxChatLogEntry[]

getUsage()

Returns token usage aggregated by (ai, model) across all steps. When a provider reports prompt-cache usage, promptTokens is the uncached input portion and cacheReadTokens / cacheCreationTokens carry the cache counters. Reset with resetUsage().

const usage = gen.getUsage(); // AxProgramUsage[]
console.log(usage[0]?.tokens?.promptTokens);
gen.resetUsage();

AxAgent and AxFlow also return flat AxChatLogEntry[] logs; composite programs set entry.name so callers can filter by node/stage.

Examples

Fetch these for full working code:

• Streaming — field-by-field streaming
• Best Of N — reward-scored sample selection
• Refine — retry rounds with generated feedback
• Streaming Assert — fail-fast partial-output correction
• Structured Output — fluent API with validation
• Debug Logging — debug mode and step hooks
• Stop Function — stop functions
• Fibonacci — streaming with thinking
• Extraction — information extraction
• Multi-Sampling — sample count usage

Native MCP/UCP

Use ax-mcp for client construction, transports, authentication, catalog and task APIs, subscriptions, event routing, and recording/replay. This section only covers the AxGen attachment boundary.

Pass live clients directly to constructor or forward options:

const gen = ax('question:string -> answer:string', { mcp: [docs, search] });
const result = await gen.forward(llm, { question }, {
  mcpContext: [
    { client: 'docs', resource: { uri: 'docs://guide' } },
  ],
});

The model receives native tool definitions. Structured, image, audio, resource-link, embedded-resource, metadata, task, and error results are preserved until the provider adapter maps supported content. Streaming keeps MCP progress/task events separate from Ax output. Never call toFunction() for native integration.

Use client.inspectCatalog() when an endpoint is the only configuration. It discovers server-owned tool/prompt names, concrete resource URIs, and URI templates. Event sources require an explicit none/all/URI/selector resource subscription policy and never create a wake route implicitly.

Under an event target, a required task-backed MCP tool registers the owning namespace:taskId continuation automatically. Use AxMCPEventSource plus axMCPEventRoutes to observe progress and resume the target on input_required or a terminal state.

Event Targets

Wrap an AxGen with eventTarget('id').program(gen).ai(ai).input(...).build() to invoke it from an explicit wake or resume route. Use segment-safe eventPath selectors; projection and explicit fields are validated against the AxGen signature before invocation. Use .wakeInput() and .resumeInput() for different action contracts. Streaming targets persist each chunk before optional chunk sinks and persist the final result before final sinks.

Use a reusable eventInput().project(...).field(...) plan when mapping should be callback-free. Callback mapInput remains available, but its result is cloned, stripped to declared AxGen inputs, and signature-validated before the first model call; mapper exceptions become non-retryable event_input_invalid deliveries.

Do Not Generate

• Do not use new AxGen(...) for new code unless explicitly required.
• Do not pass raw API keys or config objects where an ai(...) instance is expected.
• Do not use forward() for streaming; use streamingForward().
• Do not use streaming assertions as reward/refine mechanisms; they enforce hard partial-output invariants and retry with correction.
• Do not mutate step hook context expecting immediate effect; mutations are pending until the next step.
• Do not assume multi-step stops after one LLM call; it continues until outputs are filled, a stop function fires, or maxSteps is reached.