Documentation

Build LLM-powered agents
with production-ready TypeScript

DSPy for TypeScript. Working with LLMs is complex—they don't always do what you want. DSPy makes it easier to build amazing things with LLMs. Just define your inputs and outputs (signature) and an efficient prompt is auto-generated and used. Connect together various signatures to build complex systems and workflows using LLMs.

15+ LLM Providers
End-to-end Streaming
Auto Prompt Tuning

Ax, DSPy for Typescript

Working with LLMs is complex they don’t always do what you want. DSPy makes it easier to build amazing things with LLMs. Just define your inputs and outputs (signature) and an efficient prompt is auto-generated and used. Connect together various signatures to build complex systems and workflows using LLMs

And to help you really use this in production we have everything else you need like observability, streaming, support for other modalities (images,audio, etc), error-correction, multi-step function calling, MCP, RAG, etc

Screenshot 2025-06-30 at 12 52 57 PM

NPM Package Twitter Discord Chat

Why use Ax?

Production Ready

Recent Updates

v14.0.4 (Latest - Aug 2025)

v14.0.3 (Aug 2025)

v14.0.2 (Aug 2025)

v14.0.0 (Major Release - July 2025)

View Full Changelog | Migration Guide

What’s a prompt signature?

shapes at 24-03-31 00 05 55

Efficient type-safe prompts are auto-generated from a simple signature. A prompt signature is made up of a "task description" inputField:type "field description" -> "outputField:type. The idea behind prompt signatures is based on work done in the “Demonstrate-Search-Predict” paper.

You can have multiple input and output fields, and each field can be of the types string, number, boolean, date, datetime, class "class1, class2", code, json, image, audio, file, url, or an array of any of these, e.g., string[]. When a type is not defined, it defaults to string.

Field Modifiers

The recommended type-safe approach uses the ax() function with string-based signatures, which provide complete TypeScript type inference and clean, readable syntax:

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

// Create type-safe generators with ax() function
const gen1 = ax(
  "question:string -> answer:string",
);

// With field types and descriptions - fully type-safe
const gen2 = ax(
  'input:string "User input" -> category:class "tech, business, sports" "Content category", confidence:number "Confidence score 0-1"',
);

// With modifiers - clean syntax
const gen3 = ax(
  'text:string -> summary?:string "Brief summary", reasoning!:string "Internal reasoning"',
);

// TypeScript knows the exact input/output types!
type Sig2Input = { input: string };
type Sig2Output = {
  category: "tech" | "business" | "sports";
  confidence: number;
};

Type-Safe Generator Creation

Create fully typed generators using the ax() function with complete type inference:

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

// Create type-safe generators
const gen = ax("question:string -> answer:string");

// Advanced example with full type safety
const sentimentGen = ax(
  'text:string "Text to analyze" -> sentiment:class "positive, negative, neutral" "Sentiment classification", confidence:number "Confidence score 0-1"',
);

// TypeScript knows exact types - no any/unknown!
const result = await sentimentGen.forward(llm, {
  text: "I love this product!", // TypeScript validates this
});

// Perfect IntelliSense on result
console.log(result.sentiment); // "positive" | "negative" | "neutral"
console.log(result.confidence); // number

Fluent API for Complex Signatures

For building complex signatures programmatically, use the fluent API with the f() function:

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

// Build signatures step by step with fluent chaining
const sig = f()
  .input("userMessage", f.string("User input message"))
  .input("context", f.string("Background context").optional().array())
  .output("response", f.string("Generated response"))
  .output(
    "sentiment",
    f.class(["positive", "negative", "neutral"], "Sentiment"),
  )
  .description("Analyzes user messages with context")
  .build();

// Use with ax() or AxGen
const llm = ai({ name: "openai", apiKey: process.env.OPENAI_APIKEY! });
const gen = ax(sig, llm);

The fluent API supports:

Legacy: Template Literal Support (Deprecated)

Note: Template literals are deprecated. Use the ax() function instead for better type safety and performance.

Why ax() Is Fully Type-Safe

Type Safety Benefits:

Featureax() FunctionTemplate Literals (Deprecated)Raw Strings
Return TypeFully typed objectsLimited type infoany
IntelliSenseComplete field autocompletionBasic supportNo type hints
Compile-time ValidationCatches type errorsRuntime validationRuntime-only validation
Field Type InferenceExact literal typesGeneric typesNo types
IDE SupportFull refactoring supportLimitedNone
PerformanceFaster parsingTemplate overheadFastest

Example showing the recommended approach:

// ✅ Recommended: ax() function - full type safety
const recommendedWay = ax(
  'input:string -> category:class "a, b, c"',
);

// ❌ Deprecated: Template literal approach
// const deprecatedWay = ax`input:string -> category:class "a,b,c"`;

// Usage with full type safety:
const result = await recommendedWay.forward(llm, { input: "test" });
result.category; // TypeScript shows "a" | "b" | "c" - perfect safety!

Best Practices:

The ax() function provides the best balance of type safety, performance, and readability.

API Changes & Deprecations

Version 14.0.0+ introduces significant API improvements. Some older patterns are deprecated and will be removed in v15.0.0:

For detailed migration instructions and automated migration scripts, see MIGRATION.md.

Output Field Types

TypeDescriptionUsage ExampleExample Output
stringA sequence of charactersfullName:string"John Doe"
numberA numerical valueprice:number42
booleanA true or false valueisValid:booleantrue, false
dateA date valuestartDate:date"2023-10-01"
datetimeA date and time valuecreatedAt:datetime"2023-10-01T12:00:00Z"
jsonA JSON objectmetadata:json{"key": "value"}
imageAn image (input only)photo:imageBase64 encoded image data
audioAn audio file (input only)recording:audioBase64 encoded audio data
fileA file with filename, mime type, and datadocument:file{"filename": "doc.pdf", "mimeType": "application/pdf", "data": "base64data"}
urlA URL with optional title and descriptionwebsite:url"https://example.com" or {"url": "https://example.com", "title": "Example"}
class "option1,option2"Classification with predefined optionscategory:class "urgent,normal,low""urgent"
codeA code blocksolution:code "Python solution"print('Hello, world!')
string[]An array of stringstags:string[]["example1", "example2"]
number[]An array of numbersscores:number[][1, 2, 3]
boolean[]An array of boolean valuespermissions:boolean[][true, false, true]
date[]An array of datesholidayDates:date[]["2023-10-01", "2023-10-02"]
datetime[]An array of date and time valueslogTimestamps:datetime[]["2023-10-01T12:00:00Z", "2023-10-02T12:00:00Z"]
file[]An array of filesattachments:file[][{"filename": "doc1.pdf", "mimeType": "application/pdf", "data": "base64data"}]
url[]An array of URLslinks:url[]["https://example.com", {"url": "https://test.com", "title": "Test"}]
class[] "option1,option2"Array of classificationscategories:class[] "tech,business"["tech", "business"]

Important Notes on Field Types

By default, Ax enforces strict naming rules for signature fields. To allow generic names like text, input, etc., set axGlobals.signatureStrict = false. Use with caution as it may reduce signature clarity.

LLMs Supported

Google Gemini, OpenAI, OpenAI Responses, Azure OpenAI, Anthropic, X Grok, TogetherAI, Cohere, Mistral, Groq, DeepSeek, Ollama, Reka, Hugging Face

Install

Node.js / Bundlers

npm install @ax-llm/ax
# or
yarn add @ax-llm/ax

Browser (CDN)

<!-- Global variable (window.ax) -->
<script src="https://unpkg.com/@ax-llm/ax@latest?conditions=browser"></script>

<!-- Or ES modules -->
<script type="module">
  import { ax, AxAI, f } from "https://unpkg.com/@ax-llm/ax@latest";
</script>

Browser CORS Setup: Most LLM providers require a CORS proxy for browser usage. See our browser example which includes a simple CORS proxy setup.

const llm = ai({
  name: "openai",
  apiKey: "your-api-key",
  options: {
    corsProxy: "http://localhost:3001", // Your CORS proxy URL
  },
});

Example: Summarize text

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

const textToSummarize = `
The technological singularity—or simply the singularity[1]—is a hypothetical future point in time at which technological growth becomes uncontrollable and irreversible, resulting in unforeseeable changes to human civilization.[2][3] ...`;

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

const gen = ax(
  'textToSummarize:string -> textType:class "note, email, reminder", shortSummary:string "summarize in 5 to 10 words"',
);

const res = await gen.forward(llm, { textToSummarize });

console.log(">", res);

Example: Using functions for type-safe signatures

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

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

// Create a signature using string-based approach
const gen = ax(
  'userInput:string "User message or question" -> category:class "question, request, complaint" "Message type", priority:class "high, medium, low" "Urgency level", response:string "Appropriate response", reasoning!:string "Internal reasoning for classification"',
);

const res = await gen.forward(llm, {
  userInput: "My order hasn't arrived and I need it urgently!",
});

console.log("Category:", res.category);
console.log("Priority:", res.priority);
console.log("Response:", res.response);
// Note: reasoning is internal and won't appear in final output

Example: Building an agent

Use the agent prompt (framework) to build agents that work with other agents to complete tasks. Agents are easy to make with prompt signatures. Try out the agent example.

# npm run tsx ./src/examples/agent.ts

const researcher = agent({
  name: 'researcher',
  description: 'Researcher agent',
  signature: ax('physicsQuestion:string "physics questions" -> answer:string "reply in bullet points"')
});

const summarizer = agent({
  name: 'summarizer',
  description: 'Summarizer agent',
  signature: ax('text:string "text to summarize" -> shortSummary:string "summarize in 5 to 10 words"')
});

const agent = agent({
  name: 'agent',
  description: 'An agent to research complex topics',
  signature: ax('question:string -> answer:string'),
  agents: [researcher, summarizer]
});

agent.forward(llm, { questions: "How many atoms are there in the universe" })

Thinking Models Support

Ax provides native support for models with thinking capabilities, allowing you to control the thinking token budget and access the model’s thoughts. This feature helps in understanding the model’s reasoning process and optimizing token usage.

const llm = ai({
  name: "google-gemini",
  apiKey: process.env.GOOGLE_APIKEY as string,
  config: {
    model: AxAIGoogleGeminiModel.Gemini25Flash,
    thinking: { includeThoughts: true },
  },
});

// Or control thinking budget per request
const gen = ax("question:string -> answer:string");
const res = await gen.forward(
  ai,
  { question: "What is quantum entanglement?" },
  { thinkingTokenBudget: "medium" }, // 'minimal', 'low', 'medium', or 'high'
);

// Access thoughts in the response
console.log(res.thoughts); // Shows the model's reasoning process

Vector DBs Supported

Vector databases are critical to building LLM workflows. We have clean abstractions over popular vector databases and our own quick in-memory vector database.

ProviderTested
In Memory🟢 100%
Weaviate🟢 100%
Cloudflare🟡 50%
Pinecone🟡 50%
// Create embeddings from text using an LLM
const ret = await this.ai.embed({ texts: "hello world" });

// Create an in memory vector db
const db = new axDB("memory");

// Insert into vector db
await this.db.upsert({
  id: "abc",
  table: "products",
  values: ret.embeddings[0],
});

// Query for similar entries using embeddings
const matches = await this.db.query({
  table: "products",
  values: embeddings[0],
});

Alternatively you can use the AxDBManager which handles smart chunking, embedding and querying everything for you, it makes things almost too easy.

const manager = new AxDBManager({ ai, db });
await manager.insert(text);

const matches = await manager.query(
  "John von Neumann on human intelligence and singularity.",
);
console.log(matches);

RAG Documents

Using documents like PDF, DOCX, PPT, XLS, etc., with LLMs is a huge pain. We make it easy with Apache Tika, an open-source document processing engine.

Launch Apache Tika

docker run -p 9998:9998 apache/tika

Convert documents to text and embed them for retrieval using the AxDBManager, which also supports a reranker and query rewriter. Two default implementations, AxDefaultResultReranker and AxDefaultQueryRewriter, are available.

const tika = new AxApacheTika();
const text = await tika.convert("/path/to/document.pdf");

const manager = new AxDBManager({ ai, db });
await manager.insert(text);

const matches = await manager.query("Find some text");
console.log(matches);

Multi-modal DSPy

When using models like GPT-4o and Gemini that support multi-modal prompts, we support using image fields, and this works with the whole DSP pipeline.

const image = fs
  .readFileSync("./src/examples/assets/kitten.jpeg")
  .toString("base64");

const gen = ax(
  "question:string, animalImage:image -> answer:string",
);

const res = await gen.forward(llm, {
  question: "What family does this animal belong to?",
  animalImage: { mimeType: "image/jpeg", data: image },
});

When using models like gpt-4o-audio-preview that support multi-modal prompts with audio support, we support using audio fields, and this works with the whole DSP pipeline.

const audio = fs
  .readFileSync("./src/examples/assets/comment.wav")
  .toString("base64");

const gen = ax("question:string, commentAudio:audio -> answer:string");

const res = await gen.forward(llm, {
  question: "What family does this animal belong to?",
  commentAudio: { format: "wav", data: audio },
});

DSPy Chat API

Inspired by DSPy’s demonstration weaving, Ax provides AxMessage for seamless conversation history management. This allows you to build chatbots and conversational agents that maintain context across multiple turns while leveraging the full power of prompt signatures. See the example for more details.

GOOGLE_APIKEY=api-key npm run tsx ./src/examples/chat.ts
// Create a chat assistant using string-based signatures
const chatBot = ax(
  'message:string "A casual message from the user" -> reply:string "A friendly, casual response"',
);

// Start a conversation with message history
const chat: AxMessage<{ message: string }>[] = [
  { role: "user", values: { message: "Hi! How are you doing today?" } },
];

// Get first response
let response = await chatBot.forward(llm, chat);
console.log(response.reply);

// Add response to chat history
chat.push({ role: "assistant", values: { message: response.reply as string } });

// Continue conversation with context
chat.push({
  role: "user",
  values: { message: "That's great! Can you tell me a fun fact?" },
});

response = await chatBot.forward(llm, chat);
console.log(response.reply);

The conversation history is automatically woven into the prompt, allowing the model to maintain context and provide coherent responses. This works seamlessly with all Ax features including streaming, function calling, and chain-of-thought reasoning.

Streaming

Assertions

We support parsing output fields and function execution while streaming. This allows for fail-fast and error correction without waiting for the whole output, saving tokens and costs and reducing latency. Assertions are a powerful way to ensure the output matches your requirements; they also work with streaming.

// setup the prompt program
const gen = ax("startNumber:number -> next10Numbers:number[]");

// add a assertion to ensure that the number 5 is not in an output field
gen.addAssert(({ next10Numbers }: Readonly<{ next10Numbers: number[] }>) => {
  return next10Numbers ? !next10Numbers.includes(5) : undefined;
}, "Numbers 5 is not allowed");

// run the program with streaming enabled
const res = await gen.forward({ startNumber: 1 }, { stream: true });

// or run the program with end-to-end streaming
const generator = await gen.streamingForward(
  { startNumber: 1 },
  {
    stream: true,
  },
);
for await (const res of generator) {
}

The above example allows you to validate entire output fields as they are streamed in. This validation works with streaming and when not streaming and is triggered when the whole field value is available. For true validation while streaming, check out the example below. This will massively improve performance and save tokens at scale in production.

// add a assertion to ensure all lines start with a number and a dot.
gen.addStreamingAssert(
  "answerInPoints",
  (value: string) => {
    const re = /^\d+\./;

    // split the value by lines, trim each line,
    // filter out empty lines and check if all lines match the regex
    return value
      .split("\n")
      .map((x) => x.trim())
      .filter((x) => x.length > 0)
      .every((x) => re.test(x));
  },
  "Lines must start with a number and a dot. Eg: 1. This is a line.",
);

// run the program with streaming enabled
const res = await gen.forward(
  {
    question: "Provide a list of optimizations to speedup LLM inference.",
  },
  { stream: true, debug: true },
);

Field Processors

Field processors are a powerful way to process fields in a prompt. They are used to process fields in a prompt before the prompt is sent to the LLM.

const gen = ax("startNumber:number -> next10Numbers:number[]");

const streamValue = false;

const processorFunction = (value) => {
  return value.map((x) => x + 1);
};

// Add a field processor to the program
const processor = new AxFieldProcessor(
  gen,
  "next10Numbers",
  processorFunction,
  streamValue,
);

const res = await gen.forward({ startNumber: 1 });

Model Context Protocol (MCP)

Ax provides seamless integration with the Model Context Protocol (MCP), allowing your agents to access external tools, and resources through a standardized interface.

Using AxMCPClient

The AxMCPClient allows you to connect to any MCP-compatible server and use its capabilities within your Ax agents:

import { AxMCPClient, AxMCPStdioTransport } from "@ax-llm/ax-tools";

// Initialize an MCP client with a transport
const transport = new AxMCPStdioTransport({
  command: "npx",
  args: ["-y", "@modelcontextprotocol/server-memory"],
});

// Create the client with optional debug mode
const client = new AxMCPClient(transport, { debug: true });

// Initialize the connection
await client.init();

// Use the client's functions in an agent
const memoryAgent = agent({
  name: "MemoryAssistant",
  description: "An assistant with persistent memory",
  signature: "input, userId -> response",
  functions: [client], // Pass the client as a function provider
});

// Or use the client with AxGen
const memoryGen = ax("input:string, userId:string -> response:string", {
  functions: [client],
});

Using AxMCPClient with a Remote Server

Calling a remote MCP server with Ax is straightforward. For example, here’s how you can use the DeepWiki MCP server to ask questions about nearly any public GitHub repository. The DeepWiki MCP server is available at https://mcp.deepwiki.com/mcp.

import {
  AxAgent,
  AxAI,
  AxAIOpenAIModel,
  AxMCPClient,
  AxMCPStreambleHTTPTransport,
} from "@ax-llm/ax";

// 1. Initialize the MCP transport to the DeepWiki server
const transport = new AxMCPStreambleHTTPTransport(
  "https://mcp.deepwiki.com/mcp",
);

// 2. Create the MCP client
const mcpClient = new AxMCPClient(transport, { debug: false });
await mcpClient.init(); // Initialize the connection

// 3. Initialize your AI model (e.g., OpenAI)
// Ensure your OPENAI_APIKEY environment variable is set
const llm = ai({
  name: "openai",
  apiKey: process.env.OPENAI_APIKEY as string,
});

// 4. Create an AxAgent that uses the MCP client
const deepwikiAgent = agent<
  {
    // Define input types for clarity, matching a potential DeepWiki function
    questionAboutRepo: string;
    githubRepositoryUrl: string;
  },
  {
    answer: string;
  }
>({
  name: "DeepWikiQueryAgent",
  description: "Agent to query public GitHub repositories via DeepWiki MCP.",
  signature: "questionAboutRepo, githubRepositoryUrl -> answer",
  functions: [mcpClient], // Provide the MCP client to the agent
});

// 5. Formulate a question and call the agent
const result = await deepwikiAgent.forward(llm, {
  questionAboutRepo: "What is the main purpose of this library?",
  githubRepositoryUrl: "https://github.com/dosco/ax", // Example: Ax library itself
});
console.log("DeepWiki Answer:", result.answer);

This example shows how to connect to a public MCP server and use it within an Ax agent. The agent’s signature (questionAboutRepo, githubRepositoryUrl -> answer) is an assumption of how one might interact with the DeepWiki service; you would typically discover the available functions and their signatures from the MCP server itself (e.g., via an mcp.getFunctions call if supported, or documentation).

For a more complex example involving authentication and custom headers with a remote MCP server, please refer to the src/examples/mcp-client-pipedream.ts file in this repository.

Type-Safe AI Models with Automatic Inference

New in Ax: Enhanced type safety with automatic model key inference! Define your models once and get precise TypeScript types throughout your application.

Enhanced Type Inference with Static Factory Methods

Use the static .create() method for automatic type inference from your models array:

import { AxAI, AxAIGoogleGeminiModel, AxAIOpenAIModel } from "@ax-llm/ax";

// ✨ Automatic type inference on models 'fast' | 'smart' | 'reasoning'
const openai = AxAI.create({
  name: "openai",
  apiKey: process.env.OPENAI_APIKEY!,
  models: [
    {
      key: "fast" as const,
      model: AxAIOpenAIModel.GPT4OMini,
      description: "Fast model for simple tasks",
    },
    {
      key: "smart" as const,
      model: AxAIOpenAIModel.GPT4O,
      description: "Smart model for complex tasks",
    },
    {
      key: "reasoning" as const,
      model: AxAIOpenAIModel.O1Preview,
      description: "Reasoning model for deep analysis",
    },
  ],
});

// Perfect IntelliSense! The models list has exact literal types
const models = openai.getModelList();
// models[0].key is typed as 'fast' | 'smart' | 'reasoning', not just string

// Type-safe model selection in chat requests
const response = await openai.chat({
  chatPrompt: [{ role: "user", content: "Hello!" }],
  model: "fast", // ✅ TypeScript validates this is a valid key
  // model: 'invalid' // ❌ TypeScript error - not in defined models
});

Multi-Provider Type Safety

Combine multiple AI providers with precise type inference:

// Each provider gets its own inferred model keys
const gemini = AxAI.create({
  name: "google-gemini",
  apiKey: process.env.GOOGLE_APIKEY!,
  models: [
    {
      key: "quick" as const,
      model: AxAIGoogleGeminiModel.Gemini15Flash,
      description: "Quick responses",
    },
    {
      key: "advanced" as const,
      model: AxAIGoogleGeminiModel.Gemini15Pro,
      description: "Advanced reasoning",
    },
  ],
});

// MultiServiceRouter automatically infers union of all model keys
const router = new AxMultiServiceRouter([openai, gemini]);
// router now knows about 'fast' | 'smart' | 'reasoning' | 'quick' | 'advanced'

const gen = ax("inputText:string -> outputText:string");
gen.forward(router, { inputText }, { model: "quick" });
// ax now knows about 'fast' | 'smart' | 'reasoning' | 'quick' | 'advanced'

AxFlow: Build AI Workflows

AxFlow makes it easy to build complex AI workflows with automatic parallel execution and simple, readable code.

Key Features

Basic Example: Document Analysis

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

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

// Simple document analysis workflow
const documentAnalyzer = new AxFlow<
  { documentText: string },
  { summary: string; sentiment: string; keywords: string[] }
>()
  .node("summarizer", "documentText:string -> summary:string")
  .node("sentimentAnalyzer", "documentText:string -> sentiment:string")
  .node("keywordExtractor", "documentText:string -> keywords:string[]")
  // These three operations run automatically in parallel! ⚡
  .execute("summarizer", (state) => ({ documentText: state.documentText }))
  .execute(
    "sentimentAnalyzer",
    (state) => ({ documentText: state.documentText }),
  )
  .execute(
    "keywordExtractor",
    (state) => ({ documentText: state.documentText }),
  )
  // Use returns() for proper TypeScript type inference
  .returns((state) => ({
    summary: state.summarizerResult.summary,
    sentiment: state.sentimentAnalyzerResult.sentiment,
    keywords: state.keywordExtractorResult.keywords,
  }));

// Execute the workflow
const result = await documentAnalyzer.forward(llm, {
  documentText: "AI technology is revolutionary and will change everything...",
});

// TypeScript now knows the exact return type thanks to .returns()
console.log("Summary:", result.summary); // ✨ Fully typed as string
console.log("Sentiment:", result.sentiment); // ✨ Fully typed as string
console.log("Keywords:", result.keywords); // ✨ Fully typed as string[]

Compact Syntax with Aliases

For rapid development, use AxFlow’s short aliases:

// Same functionality, ultra-compact syntax
const quickAnalyzer = flow<{ text: string }>()
  .n("sum", "text:string -> summary:string") // .n() = .node()
  .n("sent", "text:string -> sentiment:string") // .n() = .node()
  .e("sum", (s) => ({ text: s.text })) // .e() = .execute()
  .e("sent", (s) => ({ text: s.text })) // .e() = .execute()
  .r((s) => ({ // .r() = .returns() for proper type inference
    result:
      `Summary: ${s.sumResult.summary}, Sentiment: ${s.sentResult.sentiment}`,
  }));

const result = await quickAnalyzer.forward(llm, {
  text: "Building the future of AI applications...",
});

Multi-Model Intelligence

Use different AI models for different tasks:

const fastAI = AxAI.create({
  name: "openai",
  apiKey: process.env.OPENAI_APIKEY!,
  models: [
    {
      key: "fast" as const,
      model: "gpt-4o-mini",
      description: "Fast responses",
    },
  ],
});

const powerAI = AxAI.create({
  name: "openai",
  apiKey: process.env.OPENAI_APIKEY!,
  models: [
    {
      key: "power" as const,
      model: "gpt-4o",
      description: "High-quality responses",
    },
  ],
});

// 🌟 The future: AI workflows that adapt and evolve with full resilience
const autonomousContentEngine = new AxFlow<
  { concept: string; targetAudience: string },
  { campaign: string }
>(
  {
    errorHandling: {
      circuitBreaker: { failureThreshold: 3, resetTimeoutMs: 60000 },
      fallbackStrategy: "graceful",
    },
    performance: {
      maxConcurrency: 4,
      adaptiveConcurrency: true,
      resourceMonitoring: { responseTimeThreshold: 10000 },
    },
  },
)
  // Neural network of specialized AI nodes
  .node(
    "conceptAnalyzer",
    "concept:string -> themes:string[], complexity:number",
  )
  .node("audienceProfiler", "audience:string -> psychographics:string")
  .node(
    "strategyArchitect",
    "themes:string[], psychographics:string -> strategy:string",
  )
  .node(
    "contentCreator",
    "strategy:string, complexity:number -> content:string",
  )
  .node("qualityOracle", "content:string -> score:number, feedback:string")
  // 🧠 These run automatically in parallel (different AI models!)
  .execute("conceptAnalyzer", (s) => ({ concept: s.concept }), {
    ai: quantumAI,
  })
  .execute("audienceProfiler", (s) => ({ audience: s.targetAudience }), {
    ai: velocityAI,
  })
  // 🎯 Strategic architecture with deep reasoning (waits for above)
  .execute("strategyArchitect", (s) => ({
    themes: s.conceptAnalyzerResult.themes,
    psychographics: s.audienceProfilerResult.psychographics,
  }), { ai: quantumAI })
  // 🎨 Creative content generation (waits for strategy)
  .execute("contentCreator", (s) => ({
    strategy: s.strategyArchitectResult.strategy,
    complexity: s.conceptAnalyzerResult.complexity,
  }), { ai: creativityAI })
  // 🔄 Quality check loop
  .label("evolve")
  .execute(
    "qualityOracle",
    (s) => ({ content: s.contentCreatorResult.content }),
    { ai: quantumAI },
  )
  .feedback((s) => s.qualityOracleResult.score < 0.9, "evolve", 3)
  // 🏆 Final transformation
  .map((s) => ({ campaign: s.contentCreatorResult.content }));

// 🚀 Execute the future
const result = await autonomousContentEngine.forward(quantumAI, {
  concept: "Sustainable AI for climate solutions",
  targetAudience: "Tech-forward environmental activists",
});

console.log("🌟 Autonomous Campaign Generated:", result.campaign);

Advanced RAG: axRAG

axRAG is a powerful, production-ready RAG (Retrieval-Augmented Generation) implementation built on AxFlow that provides advanced multi-hop retrieval, self-healing quality loops, and intelligent query refinement.

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

// Create an advanced RAG pipeline with multi-hop retrieval and self-healing
const rag = axRAG(queryVectorDB, {
  maxHops: 3, // Multi-hop context accumulation
  qualityThreshold: 0.8, // Quality-driven retrieval
  maxIterations: 2, // Parallel sub-query processing
  qualityTarget: 0.85, // Self-healing quality loops
  debug: true, // Full pipeline visualization
});

const result = await rag.forward(llm, {
  originalQuestion:
    "How do ML algorithms impact privacy in financial services?",
});

Key Features: Multi-hop retrieval, intelligent query refinement, parallel sub-query processing, self-healing quality loops, gap analysis, configurable performance vs. quality trade-offs.

For comprehensive documentation, architecture details, and advanced examples, see our detailed AxRAG Guide.

Why AxFlow is the Future

🚀 Automatic Performance Optimization:

🛡️ Production-Ready Resilience:

Compared to Traditional Approaches:

Real-World Superpowers:

“AxFlow doesn’t just execute AI workflows—it orchestrates the future of intelligent systems with automatic performance optimization”

Ready to build the impossible? AxFlow extends AxProgramWithSignature, giving you access to the entire Ax ecosystem: optimization, streaming, tracing, function calling, and more. The future of AI development is declarative, adaptive, and beautiful.

NEW: Parallel Map with Batch Size Control

Execute multiple transformations in parallel with intelligent batch processing for optimal resource management:

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

// Configure batch processing for optimal performance
const flow = new AxFlow<StateType, ResultType>({
  batchSize: 5, // Process 5 operations at a time
})
  .init({ data: largeDataset })
  // Execute multiple transforms in parallel with automatic batching
  .map([
    (state) => ({ ...state, analysis1: analyzeData(state.data) }),
    (state) => ({ ...state, analysis2: extractFeatures(state.data) }),
    (state) => ({ ...state, analysis3: generateSummary(state.data) }),
    (state) => ({ ...state, analysis4: calculateMetrics(state.data) }),
    (state) => ({ ...state, analysis5: validateResults(state.data) }),
  ], { parallel: true })
  .map((state) => ({
    final: combineResults([
      state.analysis1,
      state.analysis2,
      state.analysis3,
      state.analysis4,
      state.analysis5,
    ]),
  }));

// ⚡ Automatic batch processing: runs 5 operations concurrently,
// then processes remaining operations, maintaining result order
const result = await flow.forward(llm, { data: dataset });

🚀 Benefits:

AI Routing and Load Balancing

Ax provides two powerful ways to work with multiple AI services: a load balancer for high availability and a router for model-specific routing.

Load Balancer

The load balancer automatically distributes requests across multiple AI services based on performance and availability. If one service fails, it automatically fails over to the next available service.

import {
  AxAI,
  AxAIAnthropicModel,
  AxAIOpenAIModel,
  AxBalancer,
} from "@ax-llm/ax";

// Setup multiple AI services with specific model configurations
const openaiService = AxAI.create({
  name: "openai",
  apiKey: process.env.OPENAI_APIKEY,
  models: [
    {
      key: "smart-model",
      model: AxAIOpenAIModel.GPT4O,
      description: "Smart Model via OpenAI",
    },
    {
      key: "fast-model",
      model: AxAIOpenAIModel.GPT4OMini,
      description: "Fast Model via OpenAI",
    },
  ] as const,
});

const anthropicService = AxAI.create({
  name: "anthropic",
  apiKey: process.env.ANTHROPIC_APIKEY,
  models: [
    {
      key: "smart-model",
      model: AxAIAnthropicModel.Claude35Sonnet,
      description: "Smart Model via Anthropic",
    },
    {
      key: "fast-model",
      model: AxAIAnthropicModel.Claude35Haiku,
      description: "Fast Model via Anthropic",
    },
  ] as const,
});

// Create type-safe load balancer with automatic model key inference
// TModelKey is automatically inferred as: "smart-model" | "fast-model"
const balancer = AxBalancer.create([openaiService, anthropicService]);

// Use like a regular AI service - automatically uses the best available service
// Model key is type-safe: only "smart-model" or "fast-model" are allowed
const response = await balancer.chat({
  chatPrompt: [{ role: "user", content: "Hello!" }],
  model: "smart-model", // ✅ Type-safe
});

// Or use the balance with AxGen
const gen = ax("question:string -> answer:string");
const res = await gen.forward(balancer, { question: "Hello!" });

Multi-Service Router

The router lets you use multiple AI services through a single interface, automatically routing requests to the right service based on the model specified. With type-safe model key inference, you get automatic IntelliSense and compile-time validation.

import {
  AxAI,
  AxAIAnthropicModel,
  AxAIGoogleGeminiModel,
  AxAIOpenAIModel,
  AxMultiServiceRouter,
} from "@ax-llm/ax";

// Setup OpenAI with model list
const openaiService = AxAI.create({
  name: "openai",
  apiKey: process.env.OPENAI_APIKEY,
  models: [
    {
      key: "basic",
      model: AxAIOpenAIModel.GPT4OMini,
      description: "Model for simple tasks and quick questions",
    },
    {
      key: "medium",
      model: AxAIOpenAIModel.GPT4O,
      description: "Model for complex tasks like summarizing and coding",
    },
  ] as const,
});

// Setup Anthropic with model list
const anthropicService = AxAI.create({
  name: "anthropic",
  apiKey: process.env.ANTHROPIC_APIKEY,
  models: [
    {
      key: "deep-thinker",
      model: AxAIAnthropicModel.Claude35Sonnet,
      description: "Model for tasks requiring deep planning and analysis",
    },
  ] as const,
});

// Setup Google Gemini with model list
const googleService = AxAI.create({
  name: "google-gemini",
  apiKey: process.env.GOOGLE_APIKEY,
  models: [
    {
      key: "expert",
      model: AxAIGoogleGeminiModel.Gemini15Pro,
      description: "Model for very complex tasks and large essays",
    },
  ] as const,
});

// Create type-safe multi-service router with automatic model key inference
// TModelKey is automatically inferred as: "basic" | "medium" | "deep-thinker" | "expert"
const router = AxMultiServiceRouter.create([
  openaiService,
  anthropicService,
  googleService,
]);

// Route to specific models with full type safety
const basicResponse = await router.chat({
  chatPrompt: [{ role: "user", content: "Quick question!" }],
  model: "basic", // ✅ Routes to OpenAI GPT-4o Mini
});

const expertResponse = await router.chat({
  chatPrompt: [{ role: "user", content: "Complex analysis needed" }],
  model: "expert", // ✅ Routes to Google Gemini 1.5 Pro
});

// TypeScript will catch invalid model keys at compile time:
// model: "invalid-model" // ❌ Type error - not in union type

// Or use the router with AxGen
const gen = ax("question:string -> answer:string");
const res = await gen.forward(router, { question: "Hello!" });

🚀 Type Safety Benefits:

⚡ Use Cases:

Both classes work seamlessly with all Ax features like streaming, function calling, and chain-of-thought prompting.

OpenTelemetry support

The ability to trace and observe your llm workflow is critical to building production workflows. OpenTelemetry is an industry-standard, and we support the new gen_ai attribute namespace. Checkout src/examples/telemetry.ts for more information.

import { trace } from "@opentelemetry/api";
import {
  BasicTracerProvider,
  ConsoleSpanExporter,
  SimpleSpanProcessor,
} from "@opentelemetry/sdk-trace-base";

const provider = new BasicTracerProvider();
provider.addSpanProcessor(new SimpleSpanProcessor(new ConsoleSpanExporter()));
trace.setGlobalTracerProvider(provider);

const tracer = trace.getTracer("test");

const llm = ai({
  name: "ollama",
  config: { model: "nous-hermes2" },
  options: { tracer },
});

const gen = ax(
  'text:string -> shortSummary:string "summarize in 5 to 10 words"',
);

const res = await gen.forward({ text });
{
  "traceId": "ddc7405e9848c8c884e53b823e120845",
  "name": "Chat Request",
  "id": "d376daad21da7a3c",
  "kind": "SERVER",
  "timestamp": 1716622997025000,
  "duration": 14190456.542,
  "attributes": {
    "gen_ai.system": "Ollama",
    "gen_ai.request.model": "nous-hermes2",
    "gen_ai.request.max_tokens": 500,
    "gen_ai.request.temperature": 0.1,
    "gen_ai.request.top_p": 0.9,
    "gen_ai.request.frequency_penalty": 0.5,
    "gen_ai.request.llm_is_streaming": false,
    "http.request.method": "POST",
    "url.full": "http://localhost:11434/v1/chat/completions",
    "gen_ai.usage.completion_tokens": 160,
    "gen_ai.usage.prompt_tokens": 290
  }
}

Setting Telemetry Globally

You can set up OpenTelemetry tracing globally for all Ax operations using axGlobals. You can also pass it into AxGen, AxAI, AxAgent, etc as needed.

// Set the tracer globally for all Ax operations
axGlobals.tracer = trace.getTracer("my-app");
// Set the meter globally
axGlobals.meter = metrics.getMeter("my-app");

DSPy Metrics & Observability

Ax provides comprehensive metrics tracking for DSPy-style generation workflows, giving you deep insights into performance, error rates, and optimization opportunities. The metrics system tracks everything from basic generation latency to complex multi-step function calling patterns.

Automatic Metrics Collection

When you set up a meter globally, AxGen automatically tracks detailed metrics for all generation operations:

import { metrics } from "@opentelemetry/api";
import { axGlobals } from "@ax-llm/ax";

// Set up metrics globally
axGlobals.meter = metrics.getMeter("my-app");

// All AxGen operations now automatically track metrics
const gen = ax("userQuestion:string -> assistantAnswer:string");
const result = await gen.forward(llm, { userQuestion: "Hello!" });

Metrics Tracked

The DSPy metrics system provides comprehensive coverage of your generation workflows:

Generation Performance

Multi-step Generation

Function Calling

Streaming Performance

Performance Breakdown

Tuning the prompts (Basic)

You can tune your prompts using a larger model to help them run more efficiently and give you better results. This is done by using an optimizer like AxBootstrapFewShot with and examples from the popular HotPotQA dataset. The optimizer generates demonstrations demos which when used with the prompt help improve its efficiency.

// Download the HotPotQA dataset from huggingface
const hf = new AxHFDataLoader({
  dataset: "hotpot_qa",
  split: "train",
});

const examples = await hf.getData<{ question: string; answer: string }>({
  count: 100,
  fields: ["question", "answer"],
});

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

// Setup the program to tune
const program = ax<{ question: string }, { answer: string }>(
  'question:string -> answer:string "in short 2 or 3 words"',
);

// Setup a Bootstrap Few Shot optimizer to tune the above program
const optimize = new AxBootstrapFewShot({
  studentAI: llm,
  examples,
  options: {
    maxRounds: 3,
    maxDemos: 4,
    verboseMode: true,
  },
});

// Setup a evaluation metric em, f1 scores are a popular way measure retrieval performance.
const metricFn: AxMetricFn = ({ prediction, example }) =>
  emScore(prediction.answer as string, example.answer as string);

// Run the optimizer and remember to save the result to use later
const result = await optimize.compile(program, metricFn);

// Save the generated demos to a file
// import fs from 'fs'; // Ensure fs is imported in your actual script
fs.writeFileSync("bootstrap-demos.json", JSON.stringify(result.demos, null, 2));
console.log("Demos saved to bootstrap-demos.json");
tune-prompt ```

Tuning the prompts (Advanced, Mipro v2)

MiPRO v2 is an advanced prompt optimization framework that uses Bayesian optimization to automatically find the best instructions, demonstrations, and examples for your LLM programs. By systematically exploring different prompt configurations, MiPRO v2 helps maximize model performance without manual tuning.

Key Features

How It Works

  1. Generates various instruction candidates
  2. Bootstraps few-shot examples from your data
  3. Selects labeled examples directly from your dataset
  4. Uses Bayesian optimization to find the optimal combination
  5. Applies the best configuration to your program

Basic Usage

import { AxAI, AxChainOfThought, AxMiPRO } from "@ax-llm/ax";

// 1. Setup your AI service
const llm = ai({
  name: "google-gemini",
  apiKey: process.env.GOOGLE_APIKEY,
});

// 2. Create your program
const program = ax("input:string -> output:string");

// 3. Configure the optimizer
const optimizer = new AxMiPRO({
  studentAI: llm,
  examples: trainingData, // Your training examples
  options: {
    numTrials: 20, // Number of configurations to try
    verbose: true,
  },
});

// 4. Define your evaluation metric
const metricFn = ({ prediction, example }) => {
  return prediction.output === example.output;
};

// 5. Run the optimization
const result = await optimizer.compile(program, metricFn, {
  valset: validationData, // Optional validation set
  auto: "medium", // Optimization level
});

// 6. Use the optimized program
const result = await optimizedProgram.forward(llm, { input: "test input" });

Configuration Options

MiPRO v2 provides extensive configuration options:

OptionDescriptionDefault
numCandidatesNumber of instruction candidates to generate5
numTrialsNumber of optimization trials30
maxBootstrappedDemosMaximum number of bootstrapped demonstrations3
maxLabeledDemosMaximum number of labeled examples4
minibatchUse minibatching for faster evaluationtrue
minibatchSizeSize of evaluation minibatches25
earlyStoppingTrialsStop if no improvement after N trials5
minImprovementThresholdMinimum score improvement threshold0.01
programAwareProposerUse program structure for better proposalstrue
dataAwareProposerConsider dataset characteristicstrue
verboseShow detailed optimization progressfalse
abort-patterns.tsExample on how to abort requests

Optimization Levels

You can quickly configure optimization intensity with the auto parameter:

// Light optimization (faster, less thorough)
const result = await optimizer.compile(program, metricFn, { auto: "light" });

// Medium optimization (balanced)
const result = await optimizer.compile(program, metricFn, { auto: "medium" });

// Heavy optimization (slower, more thorough)
const result = await optimizer.compile(program, metricFn, { auto: "heavy" });

Advanced Example: Sentiment Analysis

// Create sentiment analysis program
const classifyProgram = ax<
  { productReview: string },
  { label: string }
>('productReview:string -> label:string "positive or negative"');

// Configure optimizer with advanced settings
const optimizer = new AxMiPRO({
  studentAI: llm,
  examples: trainingData,
  options: {
    numCandidates: 3,
    numTrials: 10,
    maxBootstrappedDemos: 2,
    maxLabeledDemos: 3,
    earlyStoppingTrials: 3,
    programAwareProposer: true,
    dataAwareProposer: true,
    verbose: true,
  },
});

// Run optimization and save the result
const result = await optimizer.compile(classifyProgram, metricFn, {
  valset: validationData,
});

// Save configuration for future use
const programConfig = JSON.stringify(optimizedProgram, null, 2);
await fs.promises.writeFile("./optimized-config.json", programConfig);
console.log("> Done. Optimized program config saved to optimized-config.json");

Using the Tuned Prompts

Both the basic Bootstrap Few Shot optimizer and the advanced MiPRO v2 optimizer generate demos (demonstrations) that significantly improve your program’s performance. These demos are examples that show the LLM how to properly handle similar tasks.

What are Demos?

Demos are input-output examples that get automatically included in your prompts to guide the LLM. They act as few-shot learning examples, showing the model the expected behavior for your specific task.

Loading and Using Demos

Whether you used Bootstrap Few Shot or MiPRO v2, the process of using the generated demos is the same:

import fs from "fs";
import { AxAI, AxChainOfThought, AxGen } from "@ax-llm/ax";

// 1. Setup your AI service
const llm = ai({
  name: "openai",
  apiKey: process.env.OPENAI_APIKEY,
});

// 2. Create your program (same signature as used during tuning)
const program = ax('question:string -> answer:string "in short 2 or 3 words"');

// 3. Load the demos from the saved file
const demos = JSON.parse(fs.readFileSync("bootstrap-demos.json", "utf8"));

// 4. Apply the demos to your program
program.setDemos(demos);

// 5. Use your enhanced program
const result = await program.forward(llm, {
  question: "What castle did David Gregory inherit?",
});

console.log(result); // Now performs better with the learned examples

Simple Example: Text Classification

Here’s a complete example showing how demos improve a classification task:

// Create a classification program
const classifier = ax(
  'text:string -> category:class "positive, negative, neutral"',
);

// Load demos generated from either Bootstrap or MiPRO tuning
const savedDemos = JSON.parse(
  fs.readFileSync("classification-demos.json", "utf8"),
);
classifier.setDemos(savedDemos);

// Now the classifier has learned from examples and performs better
const result = await classifier.forward(llm, {
  text: "This product exceeded my expectations!",
});

console.log(result.category); // More accurate classification

Key Benefits of Using Demos

Best Practices

  1. Save Your Demos: Always save generated demos to files for reuse
  2. Match Signatures: Use the exact same signature when loading demos
  3. Version Control: Keep demo files in version control for reproducibility
  4. Regular Updates: Re-tune periodically with new data to improve demos

Both Bootstrap Few Shot and MiPRO v2 generate demos in the same format, so you can use this same loading pattern regardless of which optimizer you used for tuning.

Complete Optimization Guide

For comprehensive documentation on optimization strategies, teacher-student architectures, cost management, and advanced techniques, see our detailed Optimization Guide.

Complete AxFlow Guide

For comprehensive documentation on building complex AI workflows, multi-model orchestration, control flow patterns, and production-ready systems, see our detailed AxFlow Guide.

Complete Telemetry Guide

For comprehensive documentation on observability, metrics, tracing, and monitoring your AI applications with OpenTelemetry integration, see our detailed Telemetry Guide.

Built-in Functions

FunctionNameDescription
Docker SandboxAxDockerSessionExecute commands within a docker environment
Embeddings AdapterAxEmbeddingAdapterFetch and pass embedding to your function
JS InterpreterAxJSInterpreterExecute JS code in a sandboxed env (Node.js only)

Check out all the examples

Use the tsx command to run the examples. It makes the node run typescript code. It also supports using an .env file to pass the AI API Keys instead of putting them in the command line.

OPENAI_APIKEY=api-key npm run tsx ./src/examples/marketing.ts
ExampleDescription
advanced-rag.ts🚀 Advanced RAG with multi-hop retrieval, self-healing quality loops, and intelligent query refinement
customer-support.tsExtract valuable details from customer communications
debug-logging.tsDebug and custom logging examples with different loggers
function.tsSimple single function calling example
food-search.tsMulti-step, multi-function calling example
result-picker.tsSelect best result from multiple field-based samples
function-result-picker.tsAdvanced result selection based on function execution
marketing.tsGenerate short effective marketing sms messages
vectordb.tsChunk, embed and search text
fibonacci.tsUse the JS code interpreter to compute fibonacci
codingWithMemory.tsCoding assistant with memory and JS interpreter (demonstrates both ax-tools features)
summarize.tsGenerate a short summary of a large block of text
template-signatures.tsType-safe signatures using AxSignature.create() (deprecated template literal examples)
rag-docs.tsConvert PDF to text and embed for rag search
react.tsUse function calling and reasoning to answer questions
agent.tsAgent framework, agents can use other agents, tools etc
streaming1.tsOutput fields validation while streaming
streaming2.tsPer output field validation while streaming
streaming3.tsEnd-to-end streaming example streamingForward()
smart-hone.tsAgent looks for dog in smart home
multi-modal.tsUse an image input along with other text inputs
balancer.tsBalance between various llm’s based on cost, etc
ax-multiservice-router.tsType-safe multi-service routing and load balancing with automatic model key inference
vertex-auth-example.tsGoogle Vertex AI authentication with dynamic API keys
docker.tsUse the docker sandbox to find files by description
prime.tsUsing field processors to process fields in a prompt
simple-classify.tsUse a simple classifier to classify stuff
mcp-client-memory.tsExample of using an MCP server for memory with Ax
mcp-client-blender.tsExample of using an MCP server for Blender with Ax
mcp-client-pipedream.tsExample of integrating with a remote MCP
tune-bootstrap.tsUse bootstrap optimizer to improve prompt efficiency
tune-mipro.tsUse mipro v2 optimizer to improve prompt efficiency
mipro-optimize.tsComplex reasoning optimization with teacher model & save
mipro-python-optimizer.tsMiPro optimization with Python service integration for advanced Bayesian optimization
mipro-chained-optimize.tsTeacher-student pipeline with cost optimization & overrides
mipro-use-optimized.tsLoad and use saved optimization with cheaper models
checkpoint-recovery.tsFault-tolerant optimization with checkpoint recovery
tune-usage.tsUse the optimized tuned prompts
telemetry.tsTrace and push traces to a Jaeger service
openai-responses.tsExample using the new OpenAI Responses API
show-thoughts.tsControl and display model reasoning thoughts
reasoning-o3-example.tsAdvanced reasoning with OpenAI o3/o4 models
use-examples.tsExample of using ‘examples’ to direct the llm
metrics-dspy.tsComprehensive DSPy metrics tracking and observability for generation workflows
optimizer-metrics.tsOptimizer metrics collection and monitoring for program tuning
ax-flow.ts🚀 Futuristic AI workflow orchestration with autonomous multi-model pipelines, adaptive loops, and self-healing agents
ax-flow-auto-parallel.ts⚡ Automatic parallelization demo - zero-config performance optimization with intelligent dependency analysis
ax-flow-enhanced-demo.ts🛡️ Production-ready AxFlow with error handling, performance optimization, and enhanced type safety features

Our Goal

Large language models (LLMs) are becoming really powerful and have reached a point where they can work as the backend for your entire product. However, there’s still a lot of complexity to manage from using the correct prompts, models, streaming, function calls, error correction, and much more. We aim to package all this complexity into a well-maintained, easy-to-use library that can work with all state-of-the-art LLMs. Additionally, we are using the latest research to add new capabilities like DSPy to the library.

How to use this library?

1. Pick an AI to work with

// Pick a LLM
const ai = new AxOpenAI({ apiKey: process.env.OPENAI_APIKEY } as AxOpenAIArgs);

2. Create a prompt signature based on your usecase

// Signature defines the inputs and outputs of your prompt program
const cot = ax("question:string -> answer:string");

3. Execute this new prompt program

// Pass in the input fields defined in the above signature
const res = await cot.forward({ question: "Are we in a simulation?" });

4. Or if you just want to directly use the LLM

const res = await ai.chat([
  { role: "system", content: "Help the customer with his questions" }
  { role: "user", content: "I'm looking for a Macbook Pro M2 With 96GB RAM?" }
]);

How do you use function calling

1. Define the functions

// define one or more functions and a function handler
const functions = [
  {
    name: "getCurrentWeather",
    description: "get the current weather for a location",
    parameters: {
      type: "object",
      properties: {
        location: {
          type: "string",
          description: "location to get weather for",
        },
        units: {
          type: "string",
          enum: ["imperial", "metric"],
          default: "imperial",
          description: "units to use",
        },
      },
      required: ["location"],
    },
    func: async (args: Readonly<{ location: string; units: string }>) => {
      return `The weather in ${args.location} is 72 degrees`;
    },
  },
];

2. Pass the functions to a prompt

const cot = ax("question:string -> answer:string", { functions });

Enable debug logs

const llm = ai({
  name: "openai",
  apiKey: process.env.OPENAI_APIKEY,
} as AxOpenAIArgs);
ai.setOptions({ debug: true });

Custom Logger

You can provide a custom logger function to control how debug information and other messages are output. This is useful for integrating with logging frameworks or customizing the output format.

import {
  AxAI,
  axCreateDefaultColorLogger,
  axCreateDefaultTextLogger,
  AxGen,
  type AxLoggerData,
} from "@ax-llm/ax";

// Custom logger that handles structured logger data
const customLogger = (data: AxLoggerData) => {
  const timestamp = new Date().toISOString();

  // Handle different types of log messages
  switch (data.name) {
    case "ChatRequestChatPrompt":
      console.log(`[${timestamp}] Chat request step ${data.step}`);
      break;
    case "ChatResponseResults":
      console.log(`[${timestamp}] Chat response: ${data.value.length} results`);
      break;
    case "FunctionResults":
      console.log(
        `[${timestamp}] Function results: ${data.value.length} calls`,
      );
      break;
    default:
      console.log(`[${timestamp}] ${data.name}:`, JSON.stringify(data.value));
  }
};

// Set logger on AI service
const llm = ai({
  name: "openai",
  apiKey: process.env.OPENAI_APIKEY,
  options: {
    debug: true,
    logger: customLogger,
  },
});

// Or set logger on generation programs
const gen = ax(
  "question:string -> answer:string",
  { logger: customLogger },
);

// Logger can also be passed through options
const result = await gen.forward(llm, { question: "Hello" }, {
  logger: customLogger,
});

Built-in Logger Factories

For convenience, Ax provides factory functions to create pre-configured loggers:

// Create a color logger that outputs to a custom function
const colorLogger = axCreateDefaultColorLogger((message: string) => {
  // Send to your logging system instead of console
  myLoggingSystem.log(message);
});

// Create a text-only logger (no colors)
const textLogger = axCreateDefaultTextLogger((message: string) => {
  fs.appendFileSync("debug.log", message + "\n");
});

const llm = ai({
  name: "openai",
  apiKey: process.env.OPENAI_APIKEY,
  options: {
    debug: true,
    logger: colorLogger,
  },
});

The logger function receives structured AxLoggerData containing different types of debug information (chat requests, responses, function calls, etc.). If no logger is provided, the default color logger is used which outputs to console.log.

Reach out

We’re happy to help reach out if you have questions or join the Discord twitter/dosco

FAQ

1. The LLM can’t find the correct function to use

Improve the function naming and description. Be very clear about what the function does. Also, ensure the function parameters have good descriptions. The descriptions can be a little short but need to be precise.

2. How do I change the configuration of the LLM I’m using?

You can pass a configuration object as the second parameter when creating a new LLM object.

const apiKey = process.env.OPENAI_APIKEY;
const conf = AxOpenAIBestConfig();
const ai = new AxOpenAI({ apiKey, conf } as AxOpenAIArgs);

3. My prompt is too long / can I change the max tokens?

const conf = axOpenAIDefaultConfig(); // or OpenAIBestOptions()
conf.maxTokens = 2000;

4. How do I change the model? (e.g., I want to use GPT4)

const conf = axOpenAIDefaultConfig(); // or OpenAIBestOptions()
conf.model = OpenAIModel.GPT4Turbo;

Monorepo tips & tricks

It is essential to remember that we should only run npm install from the root directory. This prevents the creation of nested package-lock.json files and avoids non-deduplicated node_modules.

Ask DeepWiki

Adding new dependencies in packages should be done with e.g. npm install lodash --workspace=ax (or just modify the appropriate package.json and run npm install from root).

Development Commands

# Build all workspaces
npm run build

# Run tests across all workspaces
npm run test

# Fix formatting and linting
npm run fix

# Check for circular dependencies (helps maintain clean architecture)
npm run lint:circular

# Run examples with tsx
npm run tsx ./src/examples/<example-file>.ts

# Development mode for specific workspace
npm run dev --workspace=@ax-llm/ax