Ax Agent Optimize

AxAgent Optimize Codegen Rules (@ax-llm/ax)

Use this skill for agent.optimize(...) workflows. Prefer short, modern, copyable patterns. Do not repeat general agent-authoring guidance unless the user needs it. For generic ax(...) or flow(...) tuning with top-level optimize(...), use the ax-gepa skill instead.

Your job is to help the model choose a good optimization setup for the user's actual goal:

• If the user wants better tool use, prefer action-aware tasks and either a deterministic metric or the built-in judge depending on how objective the scoring is.
• If the user wants better wording only, responder optimization may be enough.
• If the user wants reusable improvements, include artifact save/load.
• If the user wants cost, tool-use, or child-agent delegation behavior improved, make the eval tasks expose those tradeoffs explicitly.

Use These Defaults

• Use agent.optimize(...) only after the agent is already configured and runnable.
• Prefer the built-in judge path first for normal agent tuning. Most users should start with tasks that include input and criteria, then let agent.optimize(...) use its default actor target and judge-based metric.
• Keep top-level optimize(program, train, metric, options) for non-agent generators and flows; do not rewrite normal agent task-record examples to the generic helper.
• Prefer a deterministic custom metric only when success is easy to score from the prediction and task record.
• Add judgeAI plus judgeOptions when the judge should run on a stronger or separate model than the agent runtime model.
• Only reach for a plain typed AxGen evaluator when the user needs LLM-as-judge behavior outside the built-in agent.optimize(...) flow.
• Default optimize target is the actor path; do not surface target unless the user clearly wants responder-only tuning or explicit program IDs.
• Use eval-safe tools or in-memory mocks because optimization replays tasks many times.
• Prefer precise tool return schemas such as f.object(...) over vague f.json(...) whenever the agent must reason about returned fields.
• Prefer task wording with canonical entity names like "the Atlas project" instead of ambiguous labels like "Atlas" when ambiguity could trigger pointless clarification.
• Save artifacts with axSerializeOptimizedProgram(result.optimizedProgram!), then restore with axDeserializeOptimizedProgram(saved) and agent.applyOptimization(...).
• For browser-safe persistence, let the caller store the serialized JSON anywhere they want such as localStorage, IndexedDB, or a backend.
• If bootstrap is enabled, bootstrapped demos are persisted inside result.optimizedProgram.demos; raw failed traces are not saved in v1.
• Auto-promoted context fields (large undeclared inputs kept runtime-only by autoUpgrade) appear in captured traces/demos as their truncated preview string, not the full value — same as declared truncate-style contextFields. This is expected; do not treat the shortened value as a bug in the saved demos.
• For first examples, pass a plain task array instead of splitting into train and validation unless the user already has a holdout set.
• GEPA-backed agent.optimize(...) now optimizes generic components exposed by the selected target programs; target: 'actor' only tunes actor components, target: 'responder' only tunes responder components, and target: 'all' broadens the component set.
• result.optimizedProgram.componentMap is the canonical saved artifact for agent GEPA runs. It may include actor instructions, descriptions, tool descriptions/names, templates, or runtime primitives depending on what the selected target exposes.
• When child-agent delegation matters, expose the child agents as named functions and tune against realistic call/no-call tasks.

Decision Guide

Pick the optimization shape from the user's need:

• "Make the agent use tools correctly" -> keep the default actor target and use expectedActions and forbiddenActions.
• "Make final answers read better" -> consider target: 'responder', but only if the task is not mostly tool-selection or clarification behavior.
• "Make the whole agent better" -> use the default actor target first; only broaden target selection when the user clearly wants that extra scope.
• "Tune child-agent delegation" -> use tasks that exercise when to call the child agent, when to call normal tools, and when to answer directly.
• "Compare before and after" -> include a held-out task plus artifact save/load and replay.
• "Repair the tasks it keeps failing, without eroding what works" -> this is playbook territory, not GEPA: use the agent-bound playbook evolve method (in TypeScript, agent.playbook().evolve(dataset)) to mine failures into verified playbook bullets under a held-out gate. Python, Java, C++, Go, and Rust expose the same loop with native method casing; see ax-playbook. optimize(...) maximizes a metric by tuning instructions and demos; playbook evolution grows durable rules.

Choose task design carefully:

• Prefer a small number of realistic tasks over broad but vague datasets.
• Prefer concrete criteria over generic "be helpful" language.
• Prefer explicit action expectations when correctness depends on tools, recipients, dates, or side effects.
• Prefer eval-safe mocks anytime the task touches email, scheduling, external APIs, or persistence.

Make Agents Optimizable

Optimization works much better when the agent and dataset remove avoidable ambiguity:

• Prefer typed tool outputs over free-form JSON blobs so the actor can rely on exact field names.
• Tell the actor the exact tool fields it may use when payload shape matters.
• Explicitly ban invented fields if the model has any reason to guess hidden IDs or alternate key names.
• If a child agent needs parent values, declare those fields in the child signature and pass them explicitly at the call site.
• For specialist synthesis, tell the agent what narrowed context should be passed to the child agent.
• Keep maxSubAgentCalls small in examples unless the user is explicitly testing broad fan-out behavior.
• Use canonical, unambiguous task wording so the model does not burn turns asking for fake clarification.
• In JS-runtime agents, require raw runnable JavaScript only. Ban javascript: prefixes, mixed prose/code, and multi-snippet turns.

Good pattern:

• tool schema says exactly what fields exist
• task names the exact entity to look up
• actor prompt says which fields to extract before calling a child agent
• metric or judge penalizes unnecessary child-agent calls and tool misuse

Bad pattern:

• tool returns json with an underspecified shape
• task uses overloaded names like Atlas without clarifying whether that is a project, team, or account
• child agent is expected to infer hidden parent state that was never passed in its call arguments
• code agent is allowed to mix natural language with JavaScript in the same turn

Metric vs Judge

Choose the scoring path based on how objectively the task can be measured:

• Use a custom metric when you can score success directly from prediction and example.
• Use the built-in agent judge when success depends on a full-run qualitative review across tool choices, clarifications, and final output.
• Use judgeOptions.description to tell the built-in judge what to value most.
• Use helper-based judge code only when the user is not inside agent.optimize(...) and still wants LLM judging.

Quick rules:

• Tool correctness with exact expected calls or forbidden calls: prefer a deterministic metric first.
• Simple extraction or classification with known correct answers: prefer a deterministic metric.
• Open-ended assistant quality, nuanced clarification behavior, or broad synthesis quality: prefer the built-in judge.
• GEPA or optimizer flows outside agents that still need LLM judging: use a plain typed AxGen evaluator.

Important:

• A custom metric overrides the built-in judge path entirely.
• Do not introduce a dedicated judge abstraction in new examples; prefer a plain typed AxGen.
• Do not add both a custom metric and judge guidance unless the user explicitly wants two separate scoring systems and understands only the custom metric drives optimization.
• If the user builds a plain AxGen judge metric, prefer a numeric score:number output over a string tier when possible. It is simpler and less fragile in practice.

Canonical Pattern

import {
  AxAIGoogleGeminiModel,
  AxJSRuntime,
  axDefaultOptimizerLogger,
  agent,
  ai,
  f,
  fn,
  axDeserializeOptimizedProgram,
  axSerializeOptimizedProgram,
} from '@ax-llm/ax';

const tools = [
  fn('sendEmail')
    .namespace('email')
    .description('Send an email message')
    .arg('to', f.string('Recipient email address'))
    .arg('body', f.string('Email body text'))
    .returns(
      f.object({
        sent: f.boolean('Whether the email was sent'),
        to: f.string('Recipient email address'),
      })
    )
    .handler(async ({ to }) => ({ sent: true, to }))
    .build(),
];

const studentAI = ai({
  name: 'google-gemini',
  apiKey: process.env.GOOGLE_APIKEY!,
  config: { model: AxAIGoogleGeminiModel.Gemini31FlashLite, temperature: 0.2 },
});

const judgeAI = ai({
  name: 'google-gemini',
  apiKey: process.env.GOOGLE_APIKEY!,
  config: { model: AxAIGoogleGeminiModel.Gemini35Flash, temperature: 1.0 },
});

const assistant = agent('query:string -> answer:string', {
  ai: studentAI,
  judgeAI,
  contextFields: [],
  runtime: new AxJSRuntime(),
  functions: tools,
  contextPolicy: { preset: 'checkpointed', budget: 'balanced' },
  judgeOptions: {
    description: 'Prefer correct tool use over polished wording.',
    model: 'judge-model',
  },
});

const tasks = [
  {
    input: { query: 'Send an email to Jim saying good morning.' },
    criteria: 'Use the email tool and send the message to Jim.',
    expectedActions: ['email.sendEmail'],
  },
];

const result = await assistant.optimize(tasks, {
  maxMetricCalls: 12,
  verbose: true,
});

const saved = axSerializeOptimizedProgram(result.optimizedProgram!);
const restored = axDeserializeOptimizedProgram(saved);
assistant.applyOptimization(restored);

Minimal Normal-User Pattern

Start here unless the user clearly needs a hand-built scorer:

const tasks = [
  {
    input: { query: 'Send an email to Jim saying good morning.' },
    criteria: 'Use the email tool and send the message to Jim.',
    expectedActions: ['email.sendEmail'],
  },
];

const result = await assistant.optimize(tasks);
assistant.applyOptimization(result.optimizedProgram!);

• target defaults to actor optimization.
• metric defaults to the built-in LLM judge.
• judgeAI is optional; if omitted, the agent falls back to its configured judge model or runtime model.
• bootstrap: true is a good next step for tool-heavy agents when you want GEPA to start from successful traces from the provided tasks.
• The one thing users still need is realistic task records with clear criteria.