Ax from 1 to 100

1. Mental Model

Ax is easiest to understand as a compiler and runtime for typed AI programs. You describe what data goes in and what data should come out. Ax turns that contract into a provider-specific prompt, streams and parses the model response, validates it, retries when possible, and returns a typed TypeScript object.

2. Setup

Install the core package and create an AI service with the provider you want. The same Ax program can usually switch providers by changing the service config.

npm install @ax-llm/ax

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

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

const classify = ax('review:string -> sentiment:class "positive, negative, neutral"');

const result = await classify.forward(llm, {
  review: "This product is amazing!",
});

console.log(result.sentiment); // typed as "positive" | "negative" | "neutral"

3. Signatures

Signatures are Ax's central abstraction. They define input fields, output fields, types, optionality, validation constraints, and sometimes prompt hints. If you understand signatures, the rest of Ax becomes composition.

Syntax

[description] input1:type, input2:type -> output1:type, output2:type

Common Types

stringnumberbooleanjsondatedatetimeurlcodeclassimageaudiofile

import { ax, f } from "@ax-llm/ax";

const extractor = ax(`
  customerEmail:string, currentDate:datetime ->
  priority:class "high, normal, low",
  sentiment:class "positive, negative, neutral",
  ticketNumber?:number,
  nextSteps:string[]
`);

const productExtractor = ax(
  f()
    .input("productPage", f.string("Raw product page text"))
    .output("product", f.object({
      name: f.string(),
      price: f.number(),
      specs: f.object({
        materials: f.string().array(),
      }),
    }))
    .build(),
);

4. AxGen

ax(...) creates an AxGen: a single-step AI program. It is the basic unit for extraction, classification, rewriting, scoring, summarization, multimodal analysis, and tool-using ReAct-style calls.

const qa = ax('userQuestion:string -> answer:string, confidence:number');

const { answer, confidence } = await qa.forward(llm, {
  userQuestion: "What is TypeScript?",
});

Generation Lifecycle

1. Render a prompt from the signature and inputs.
2. Call the selected AI provider.
3. Stream and parse output fields.
4. Validate field values and schema constraints.
5. Retry with correction feedback when validation fails.
6. Return the typed final object.

5. Providers

Ax exposes provider services through ai({ name }). The README lists OpenAI, Anthropic, Google Gemini, Mistral, Cohere, Groq, Together, Ollama, OpenRouter, Bedrock via optional package, Reka, DeepSeek, Grok, HuggingFace, and WebLLM.

const openai = ai({ name: "openai", apiKey: process.env.OPENAI_APIKEY! });
const anthropic = ai({ name: "anthropic", apiKey: process.env.ANTHROPIC_APIKEY! });
const gemini = ai({ name: "google-gemini", apiKey: process.env.GOOGLE_APIKEY! });

// Same program, different runtime service.
await classify.forward(openai, { review });
await classify.forward(anthropic, { review });
await classify.forward(gemini, { review });

6. Streaming, Validation, and Retries

Ax is streaming-first. Even when you call forward() and only need a final object, Ax can parse fields as they arrive. With streamingForward(), callers can consume output incrementally.

const writer = ax('topic:string -> title:string, outline:string[], article:string');

for await (const chunk of writer.streamingForward(llm, { topic: "AI evals" })) {
  console.log(chunk);
}

Validation constraints such as min, max, email, url, regex, and enum/class choices can feed the retry pipeline. This is one of Ax's biggest practical advantages over raw provider SDK calls.

7. Tools and Function Calling

Ax programs can call host functions. For modern code, define tools with fn() for typed arguments, typed returns, examples, namespaces, and schema validation.

import { ax, f, fn } from "@ax-llm/ax";

const searchDocs = fn("searchDocs")
  .description("Search indexed documentation")
  .namespace("kb")
  .arg("query", f.string("Search query"))
  .arg("limit", f.number("Max results").optional())
  .returns(f.string("Matching snippets").array())
  .handler(async ({ query, limit }) => {
    return mySearchIndex.search(query, { limit: limit ?? 5 });
  })
  .build();

const assistant = ax("question:string -> answer:string", {
  functions: [searchDocs],
});

Use tool calls for things the model should not invent: current data, private state, database lookups, calculations, side effects, and calls to specialist systems.

8. AxAgent

agent(...) turns a signature into a long-running, tool-using actor. The README describes a three-stage pipeline: distiller, executor, responder. The executor can run an RLM loop with a sandboxed JavaScript runtime, tools, child agents, memories, and skills.

import { agent, AxJSRuntime } from "@ax-llm/ax";

const analyzer = agent(
  "context:string, query:string -> answer:string, evidence:string[]",
  {
    agentIdentity: {
      name: "Document Analyzer",
      description: "Analyzes long documents with iterative sub-queries",
    },
    contextFields: ["context"],
    runtime: new AxJSRuntime(),
    maxTurns: 20,
    contextPolicy: { preset: "checkpointed", budget: "balanced" },
  },
);

const result = await analyzer.forward(llm, {
  context: veryLongDocument,
  query: "What are the main arguments and evidence?",
});

When to use an agent

• The task needs multiple tool calls or branching decisions.
• The model must inspect large context gradually.
• You want child agents with distinct responsibilities.
• You need clarification, runtime state, memories, or dynamic skills.

When not to use an agent

• A single extraction, classification, or rewrite is enough.
• The task has a deterministic pipeline better represented as AxFlow.
• You do not have a clear tool boundary or stopping condition.

9. AxFlow

flow() is a typed workflow builder. Define nodes, map state into node inputs, let independent nodes auto-parallelize, then return a final typed output.

import { flow } from "@ax-llm/ax";

const emailFlow = flow<{ emailText: string }, { priority: string; rationale: string }>()
  .description("Email Priority", "Classify priority and write a short rationale")
  .node("classifier", 'emailText:string -> priority:class "high, normal, low"')
  .node("rationale", "emailText:string, priority:string -> rationale:string")
  .execute("classifier", (state) => ({ emailText: state.emailText }))
  .execute("rationale", (state) => ({
    emailText: state.emailText,
    priority: state.classifierResult.priority,
  }))
  .returns((state) => ({
    priority: state.classifierResult.priority,
    rationale: state.rationaleResult.rationale,
  }));

Use flows when the structure of the process is known: classify then route, parallel extractors then merge, iterative improve-and-check loops, map-reduce over arrays, or multi-model pipelines.

10. Optimization

Ax follows DSPy's thesis: AI programs should be optimized from examples and metrics, not endlessly hand-prompted. Ax includes optimizers such as GEPA, ACE, and bootstrap few-shot.

const optimizer = new AxGEPA({
  studentAI,
  teacherAI,
  numTrials: 16,
  minibatch: true,
  minibatchSize: 6,
  seed: 42,
});

const result = await optimizer.compile(
  emailFlow,
  trainSet,
  async ({ prediction, example }) => ({
    accuracy: prediction.priority === example.priority ? 1 : 0,
    brevity: prediction.rationale.length <= 60 ? 1 : 0.4,
  }),
  { auto: "medium", validationExamples: valSet, maxMetricCalls: 240 },
);

11. Advanced Topics

12. The 1-100 Roadmap

13. Pitfalls and Sharp Edges

14. Resources and Notes

Sources used for these notes: the GitHub README, the package README, Ax signature reference, AxAgent reference, AxFlow reference, and DeepWiki's architecture overview for ax-llm/ax.

• GitHub repository
• Main README
• Signature reference
• Agent reference
• Flow reference
• DeepWiki docs