Making your agent optimizable¶
Optimizers can only tune the parts of your agent that you expose to them. This page shows how to declare that surface in your Koog strategy so the optimizers have something to work on.
What gets optimized¶
An optimizer sees your agent as a set of tunable modules — the same idea as a DSPy module: a named unit with an instruction (and optional demonstrations) that the optimizer is free to rewrite. There are two levels:
-
The strategy system prompt — the
__strategy__module. This is the non-empty system prompt you pass when you build the agent. If the agent has no system prompt, or its system prompt is explicitly empty, no strategy-level module is discovered; optimizers can still tune any optimizable subgraphs you declare. -
Optimizable subgraphs — finer-grained, per-step modules inside your strategy (the closest analogue to a DSPy module). When a single system prompt is too coarse — your agent reasons in stages, and each stage needs its own focused instruction — you declare those stages as optimizable subgraphs. The optimizer then tunes each one independently.
GEPA, MIPROv2, and BootstrapFewShot optimize the strategy system prompt, when present and
non-empty, plus every optimizable subgraph you declare. ACE currently optimizes only the
strategy-level context (it curates a playbook for the __strategy__ module); per-subgraph
optimization for ACE is a work in progress.
You don't need subgraphs to start. Reach for them when one prompt can't capture everything the agent must get right.
optimizableSubgraphWithTask¶
Declare an optimizable subgraph inside a strategy { } block with property delegation (by).
It is the optimization-aware counterpart of Koog's subgraphWithTask: a subgraph that runs a
single task whose instruction and few-shot demonstrations an optimizer can replace without
you touching the graph.
import ai.koog.agents.optimization.core.optimizableSubgraphWithTask
val analyze by optimizableSubgraphWithTask<String, String>(
optimizableInstruction = "Look at the customer message.",
freshHistory = true,
) { instruction, message -> "Customer message:\n$message" }
The key parameters:
optimizableInstruction— the starting instruction. This is what runs before any optimization; an optimizer may replace it with a better one. Start vague or start with your best guess — either way the optimizer evolves it from here.name— the module's name, defaulting to the delegated property name (analyzeabove). It must be unique within the strategy: the optimization artifact uses it as a lookup key.freshHistory— whentrue, the subgraph drops the conversation turns that came before it but keeps the agent's system messages, then appends the resolved instruction as one more system message, with the task query after it. It lets you (a) isolate a subgraph from prior chatter, (b) give each subgraph its own independently optimizable instruction on top of the shared system prompt, and (c) keep prompts focused and cheaper for classification/routing-style steps. Note the consequence: a fresh subgraph normally runs with two system messages — the agent's and its own — and optimizing the agent-level prompt changes the context every subgraph sees.defineTask { instruction, input -> ... }— builds the user message for this subgraph from the (resolved) instruction and the subgraph input, returning aString. Theinstructionargument already reflects any optimization; you usually compose your query frominputand let the system message carry the instruction.
Few-shot knobs fewShotPromptType and demonstrationFormat control how learned demonstrations are
inserted into the prompt (as message history vs. a single string, compact vs. detailed); leave them
null to inherit sensible defaults. As a brief aside, the tool/LLM knobs — toolSelectionStrategy,
llmModel, llmParams, runMode, assistantResponseRepeatMax, responseProcessor — let a
subgraph use its own tools or model; they behave exactly as in a regular subgraphWithTask.
A worked example: a support-ticket router¶
This mirrors the runnable OptimizeSupportRouterExample. The agent is a two-stage router with two
optimizable subgraphs:
Both subgraphs start from deliberately vague instructions that say nothing about the allowed routing codes or the required output format.
import ai.koog.agents.core.agent.entity.AIAgentGraphStrategy
import ai.koog.agents.core.agent.entity.createStorageKey
import ai.koog.agents.core.dsl.builder.node
import ai.koog.agents.core.dsl.builder.strategy
import ai.koog.agents.optimization.core.optimizableSubgraphWithTask
private val origMessageKey = createStorageKey<String>("orig-support-message")
private fun supportRouterStrategy(): AIAgentGraphStrategy<String, String> =
strategy("support-router") {
val remember by node<String, String>("remember") { message ->
storage.set(origMessageKey, message)
message
}
// Optimizable subgraph #1: ambiguous; doesn't say what to extract or why.
val analyze by optimizableSubgraphWithTask<String, String>(
optimizableInstruction = "Look at the customer message.",
freshHistory = true,
) { _, message -> "Customer message:\n$message" }
val pack by node<String, String>("pack") { analysis ->
val message = storage.getValue(origMessageKey)
"Customer message:\n$message\n\nNotes:\n$analysis"
}
// Optimizable subgraph #2: ambiguous; no codes, no output-format requirement.
val route by optimizableSubgraphWithTask<String, String>(
optimizableInstruction = "Decide how this request should be handled.",
freshHistory = true,
) { _, packed -> packed }
edge(nodeStart forwardTo remember)
edge(remember forwardTo analyze)
edge(analyze forwardTo pack)
edge(pack forwardTo route)
edge(route forwardTo nodeFinish)
}
The dataset labels each ticket with the routing code it should produce, and the metric is an exact match on that code:
val dataset: TrainSet<String, String> = listOf(
TrainSetItem("I was charged twice for my subscription this month.", "BILLING"),
TrainSetItem("The app crashes every time I open the settings page.", "TECH"),
TrainSetItem("Do you have a referral program for existing customers?", "OTHER"),
// ...
)
val metric: (TrainSetItem<String, String>, String) -> Double = { item, answer ->
if (answer.trim() == item.itemLabel) 1.0 else 0.0
}
The payoff. With the vague baseline instructions the agent has no idea the answer must be a bare
code, so it replies with prose like "This should go to billing." and scores ~0 on the exact-match
metric. An optimizer — GEPA here — reflects on the failure traces (which carry the expected labels),
and evolves both subgraph instructions until analyze extracts the right signals and route
emits exactly BILLING / TECH / OTHER. The learned instructions are saved to the optimization
artifact and reapplied when you load the optimized agent.
You wire the optimizer to this agent exactly as in getting-started.md: build a
trainingSession, call optimizer.train(session), then optimizer.loadOptimizedAgent(agent). See
optimizers.md for how to pick and configure GEPA (or any other optimizer); we don't
repeat that here.
Tips and gotchas¶
- Names are artifact keys. Each optimizable subgraph must have a unique name within its
strategy. The artifact maps instructions and demonstrations to subgraphs by name, so duplicates
silently cause the wrong instruction to land on the wrong subgraph. Let names default to the
property name and keep those distinct, or pass an explicit
name. freshHistory = truesuits classification-style subtasks. When a subgraph should decide something from its input alone — classify, route, extract — a fresh history keeps it focused and makes its prompt cleanly optimizable. Leave itfalsewhen the step genuinely needs the prior conversation.- The two levels compose. Strategy-level and subgraph-level optimization are not exclusive —
when the strategy system prompt is non-empty, an optimizer tunes the
__strategy__module and every optimizable subgraph in the same run. Start with just the strategy prompt; add optimizable subgraphs as you find steps that need their own instructions.
Next¶
- optimizers.md — choose and configure an optimizer (GEPA, MIPROv2, BootstrapFewShot, ACE).
- getting-started.md — the minimal end-to-end loop: build agent, dataset,
metric,
trainingSession, train, and run the optimized agent.