261 lines
8.7 KiB
Markdown
261 lines
8.7 KiB
Markdown
---
|
|
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.33"
|
|
---
|
|
|
|
# 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 by default.
|
|
- Avoid extra LLM calls inside the metric unless the user explicitly wants judge-based evaluation.
|
|
- If the user needs LLM-as-judge scoring for a non-agent GEPA run, prefer a plain typed `AxGen` evaluator instead of writing a custom judge abstraction.
|
|
- `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.
|
|
|
|
## Metric Selection
|
|
|
|
Choose the evaluation path deliberately:
|
|
|
|
- Prefer a deterministic metric when correctness can be read directly from `prediction` and `example`.
|
|
- Prefer a deterministic metric when cost, latency, recursion depth, or tool count matters.
|
|
- Use a plain typed `AxGen` evaluator only when the task is genuinely qualitative and hard to score exactly.
|
|
- For `agent.optimize(...)`, prefer the built-in judge path instead of manually wrapping a judge metric.
|
|
|
|
Rule of thumb:
|
|
|
|
- `AxGEPA` on `AxGen` or flow: use a metric first, optionally a plain typed `AxGen` evaluator if needed.
|
|
- `agent.optimize(...)`: use custom `metric` for crisp scoring, otherwise `judgeAI` plus `judgeOptions`.
|
|
|
|
## Canonical Scalar Pattern
|
|
|
|
```typescript
|
|
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
|
|
|
|
```typescript
|
|
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
|
|
|
|
```typescript
|
|
// 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
|
|
|
|
```typescript
|
|
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
|
|
|
|
```typescript
|
|
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`
|