ax-gepa

name	ax-gepa
description	This skill helps an LLM generate correct AxGEPA optimization code using @ax-llm/ax. Use when the user asks about AxGEPA, GEPA, Pareto optimization, multi-objective prompt tuning, reflective prompt evolution, validationExamples, maxMetricCalls, or optimizing a generator, flow, or agent tree.
version	19.0.17

AxGEPA Codegen Rules (@ax-llm/ax)

Use this skill to generate direct AxGEPA optimization code. Prefer short, modern, copyable patterns over long explanation.

Use These Defaults

• Use new AxGEPA({ studentAI, teacherAI, ... }).
• Prefer ai(), ax(), and flow() for new code.
• Use a strong teacherAI and a cheaper studentAI.
• Always pass validationExamples to compile().
• Always set maxMetricCalls to bound optimizer cost.
• Use scalar metrics for one objective and object metrics for Pareto optimization.
• Apply results with program.applyOptimization(result.optimizedProgram!).
• For tree-wide runs, expect optimizedProgram.instructionMap.

Critical Rules

• AxGEPA.compile() works for a single generator and for tree-aware roots such as flows or agents with registered instruction-bearing descendants.
• There is no separate flow-only GEPA optimizer. Use AxGEPA for flows too.
• The metric may return either number or Record<string, number>.
• Keep metrics deterministic and cheap. Avoid extra LLM calls inside the metric unless the user explicitly wants judge-based evaluation.
• maxMetricCalls must be large enough to cover the initial validation pass over validationExamples.
• GEPA optimizes instructions. If a tree has no instruction-bearing nodes, optimization will fail.
• Use held-out validation examples for selection. Do not reuse the training set as validationExamples.
• result.optimizedProgram is the easy-to-apply best candidate. result.paretoFront is the full trade-off set for multi-objective runs.

Canonical Scalar Pattern

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

const student = ai({
  name: 'openai',
  apiKey: process.env.OPENAI_APIKEY!,
  config: { model: AxAIOpenAIModel.GPT4OMini },
});

const teacher = ai({
  name: 'openai',
  apiKey: process.env.OPENAI_APIKEY!,
  config: { model: AxAIOpenAIModel.GPT4O },
});

const classifier = ax(
  'emailText:string -> priority:class "high, normal, low", rationale:string'
);

const train = [
  { emailText: 'URGENT: Server down!', priority: 'high' },
  { emailText: 'Weekly newsletter', priority: 'low' },
];

const validation = [
  { emailText: 'Invoice overdue', priority: 'high' },
  { emailText: 'Lunch plans?', priority: 'low' },
];

const metric = ({ prediction, example }: { prediction: any; example: any }) =>
  prediction?.priority === example?.priority ? 1 : 0;

const optimizer = new AxGEPA({
  studentAI: student,
  teacherAI: teacher,
  numTrials: 12,
  minibatch: true,
  minibatchSize: 4,
  earlyStoppingTrials: 4,
  sampleCount: 1,
});

const result = await optimizer.compile(classifier, train, metric, {
  validationExamples: validation,
  maxMetricCalls: 120,
});

classifier.applyOptimization(result.optimizedProgram!);
console.log(result.bestScore);

Canonical Pareto Pattern

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

const student = ai({
  name: 'openai',
  apiKey: process.env.OPENAI_APIKEY!,
  config: { model: AxAIOpenAIModel.GPT4OMini },
});

const teacher = ai({
  name: 'openai',
  apiKey: process.env.OPENAI_APIKEY!,
  config: { model: AxAIOpenAIModel.GPT4O },
});

const wf = flow<{ emailText: string }>()
  .n('classifier', 'emailText:string -> priority:class "high, normal, low"')
  .n(
    'rationale',
    'emailText:string, priority:string -> rationale:string "One concise sentence"'
  )
  .e('classifier', (state) => ({ emailText: state.emailText }))
  .e('rationale', (state) => ({
    emailText: state.emailText,
    priority: state.classifierResult.priority,
  }))
  .r((state) => ({
    priority: state.classifierResult.priority,
    rationale: state.rationaleResult.rationale,
  }));

const train = [
  { emailText: 'URGENT: Server down!', priority: 'high' },
  { emailText: 'Weekly newsletter', priority: 'low' },
];

const validation = [
  { emailText: 'Invoice overdue', priority: 'high' },
  { emailText: 'Lunch plans?', priority: 'low' },
];

const metric = ({ prediction, example }: { prediction: any; example: any }) => {
  const accuracy = prediction?.priority === example?.priority ? 1 : 0;
  const rationale = typeof prediction?.rationale === 'string'
    ? prediction.rationale
    : '';
  const brevity = rationale.length <= 40 ? 1 : rationale.length <= 80 ? 0.5 : 0.1;
  return { accuracy, brevity };
};

const result = await new AxGEPA({
  studentAI: student,
  teacherAI: teacher,
  numTrials: 16,
  minibatch: true,
  minibatchSize: 6,
  earlyStoppingTrials: 5,
  sampleCount: 1,
}).compile(wf, train, metric, {
  validationExamples: validation,
  maxMetricCalls: 240,
});

for (const point of result.paretoFront) {
  console.log(point.scores, point.configuration);
}

wf.applyOptimization(result.optimizedProgram!);
console.log(result.optimizedProgram?.instructionMap);

Metric Patterns

// Scalar objective
const scalarMetric = ({ prediction, example }) =>
  prediction.answer === example.answer ? 1 : 0;

// Multi-objective
const multiMetric = ({ prediction, example }) => ({
  accuracy: prediction.answer === example.answer ? 1 : 0,
  brevity:
    typeof prediction?.reasoning === 'string' &&
    prediction.reasoning.length < 120
      ? 1
      : 0.2,
});

• Return plain numbers or plain object literals.
• Keep objective names stable across calls.
• Prefer normalized scores such as 0..1 so trade-offs are easy to reason about.

Result Handling

const { optimizedProgram, paretoFront } = result;

program.applyOptimization(optimizedProgram!);

// Save for later
const saved = JSON.stringify(optimizedProgram);

// Load later and re-apply
const loaded = JSON.parse(saved);
program.applyOptimization(loaded);

• Single-target runs usually populate both optimizedProgram.instruction and optimizedProgram.instructionMap.
• Tree-wide runs rely on instructionMap, keyed by full program ID.
• Pareto points expose candidate configs under point.configuration.instructionMap.

Useful Options

const optimizer = new AxGEPA({
  studentAI,
  teacherAI,
  numTrials: 20,
  minibatch: true,
  minibatchSize: 5,
  minibatchFullEvalSteps: 5,
  earlyStoppingTrials: 5,
  minImprovementThreshold: 0,
  sampleCount: 1,
  seed: 42,
  verbose: true,
});

• numTrials: number of reflection/evolution rounds.
• minibatch: reduce per-round evaluation cost.
• minibatchSize: examples per minibatch.
• earlyStoppingTrials: stop after repeated non-improvement.
• minImprovementThreshold: reject tiny gains below this threshold.
• seed: stabilize sampling during demos and tests.

Budgeting and Validation

• Always create distinct train and validationExamples arrays.
• Size maxMetricCalls for at least one full validation pass plus several rounds.
• If the user wants a strict budget, say so explicitly and set maxMetricCalls.
• For expensive trees, start with auto: 'light' or fewer numTrials, then scale up.

Troubleshooting

• Error about maxMetricCalls being too small: increase it until the initial validation pass fits.
• Empty or poor Pareto front: verify the metric returns numbers for every example.
• No tree optimization effect: ensure child programs are registered under the root and have instructions to mutate.
• Saved optimization applies only partly: use program.applyOptimization(...), not just setInstruction(...), so instructionMap reaches the full tree.

Good Example Targets

• /Users/vr/src/ax/src/examples/gepa.ts
• /Users/vr/src/ax/src/examples/gepa-flow.ts
• /Users/vr/src/ax/src/examples/gepa-train-inference.ts
• /Users/vr/src/ax/src/examples/gepa-quality-vs-speed-optimization.ts