AxLearn: Self-Improving Agents

AxLearn provides a zero-configuration optimization loop that enables your Ax agents to automatically improve their prompts using production logs and teacher models.

Quick Start

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

// 1. Define storage (simple functional interface)
//    In a real app, save to a database.
const traces = new Map();
const storage = {
  save: async (name, item) => name === 'traces' ? traces.set(item.id, item) : null,
  load: async (name, query) => Array.from(traces.values())
};

// 2. Create your self-improving agent
//    AxLearn handles logging, trace collection, and optimization automatically.
const agent = new AxLearn(ax('customer_query -> polite_reply'), {
  name: 'support-chat-v1', // Unique identifier for versioning
  storage,                 // Where to save traces
  // Teacher for optimization
  teacher: ai({
    name: 'openai', 
    apiKey: process.env.OPENAI_APIKEY as string, 
    config: { model: AxAIOpenAIModel.GPT4O } 
  }), 
  budget: 20               // Max optimization rounds
});

// 3. Use in production
//    Traces are automatically logged to your storage.
const response = await agent.forward(ai, { customer_query: 'Where is my order?' });

// 4. Optimize offline
//    Uses collected traces + synthetic data to improve the prompt.
await agent.optimize();

How It Works

AxLearn is an all-in-one wrapper that manages the lifecycle of a self-improving agent.

1. Production Logging: When you call agent.forward(), it automatically logs input/output traces to your provided storage.
2. Data Management: During optimization, it loads these traces and mixes them with synthetic data generated by the teacher model.
3. Automatic Tuning: It runs an optimization loop (using GEPA) to iteratively improve the prompt and examples.
4. Evaluation: It uses an LLM-as-a-Judge to evaluate performance against the teacher model.

Core Components

AxStorage (Required)

You must provide a simple storage adapter to save traces and checkpoints.

Interface:

type AxStorage = {
  save: (name: string, item: AxTrace | AxCheckpoint) => Promise<void>;
  load: (name: string, query: AxStorageQuery) => Promise<(AxTrace | AxCheckpoint)[]>;
};

Example (In-Memory):

const traces = new Map<string, AxTrace[]>();

const storage: AxStorage = {
  save: async (name, item) => {
    const list = traces.get(name) ?? [];
    list.push(item);
    traces.set(name, list);
  },
  load: async (name, query) => {
    return traces.get(name) ?? [];
  }
};

AxLearn Configuration

const agent = new AxLearn(generator, {
  name: 'my-agent',
  storage: myStorage,
  
  // Optimizer Settings
  teacher: myTeacherLLM,    // Strong model for teaching/judging
  budget: 20,               // Max optimization steps
  
  // Data Settings
  useTraces: true,          // Use production traces? (default: true)
  generateExamples: true,   // Create synthetic data? (default: true)
  synthCount: 50,           // How many synthetic examples to generate
});

Workflow

┌──────────────────────────────────────────────────────────────┐
│  Production Runtime                                          │
│  ┌─────────┐                                ┌──────────────┐ │
│  │ AxLearn │──(automatic logging)──────────▶│ AxStorage    │ │
│  └─────────┘                                └──────────────┘ │
└──────────────────────────────────────────────────────────────┘
       │                                              │
       │ (agent.forward)                              │ (load traces)
       ▼                                              ▼
┌──────────────────────────────────────────────────────────────┐
│  Offline Optimization (agent.optimize)                       │
│                                                              │
│  1. Load Traces from Storage                                 │
│  2. Generate Synthetic Data (using Teacher)                  │
│  3. Run Optimization Loop (adjust prompt/examples)           │
│  4. Evaluate (using Teacher as Judge)                        │
│  5. Save Improved Checkpoint                                 │
└──────────────────────────────────────────────────────────────┘

Best Practices

1. Start with traces: Deploy your agent early to collect real-world data.
2. Use user feedback: Save traces with feedback to guide optimization.

   // Add feedback to a trace
   trace.feedback = { score: 1, label: 'good' };
   await storage.save('my-agent', trace);

3. Optimize offline: Run agent.optimize() periodically (e.g., nightly).
4. Versioning: Change the name (e.g., 'agent-v2') when you deploy a major update to keep traces separate.