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