Prompts That Can Write Themselves: A Kotlin Experiment

TL;DR

Koog Agents Optimization Library is a research library developed by JetBrains Research for optimizing the prompts of Koog agents. With it, you write a first version of each instruction, mark the steps worth tuning, provide the library with a training set, and specify a way to score the agent’s output. The library then runs your agent, learns from what happens, and hands back better instructions as a separate JSON file you can read, review, and commit. No rewrite of the agent is needed.

The examples in this post use GEPA, a prompt optimizer that learns from your agent’s failures and successes. In this post, we walk you through examples running on the privacy benchmark PUPA. GEPA improved our agent’s score from 81.19 to 94.22.

👉 Koog Agent Optimization Library (GitHub)

👉Docs

👉 Examples

An agent to optimize

Koog is JetBrains’ Kotlin framework for building agents. It treats an agent as a program: you declare control flow as a typed graph of nodes and edges, and each step can carry its own prompt, tools, and model.

We have a whole agent written in it: our port of PAPILLON, a privacy pipeline that GEPA uses as one of its benchmarks. It runs on PUPA, a dataset of user queries that carry personal details. The agent has to answer such a query with help from a big external model without handing that model any of those details, so it rewrites the question without the details, sends that version out, and writes the reply from what comes back.

fun pupaStrategy(untrustedModel: LLModel) = strategy("pupa-papillon") {
    val craftRedactedRequest by subgraphWithTask(
        name = "craft_redacted_request",
    ) { input -> REDACT_INSTRUCTION + "\n\nuser_query:\n${input.userQuery}" }

    val callUntrustedModel by subgraph(
        name = "untrusted_model",
        llmModel = untrustedModel,
    ) { /* one call to the big model */ }

    val respondToQuery by subgraphWithTask(
        name = "respond_to_query",
    ) { input -> RESPOND_INSTRUCTION + "\n\n" + input.rendered() }

    // …plus three small nodes carrying state
}

Only three steps are needed to do the work here: the first rewrites the question, the second calls the big model, and the third writes the reply. Everything else is plumbing. Two of those three steps hold an instruction that a human wrote, and they are the ones you would want an optimizer to take over.

Marking a step as optimizable

The library adds one builder. Swap subgraphWithTask for optimizableSubgraphWithTask on the steps you want tuned, and lift the instruction out of the string it was buried in:

fun pupaStrategy(untrustedModel: LLModel) = strategy("pupa-papillon") {
    val craftRedactedRequest by optimizableSubgraphWithTask(
        name = "craft_redacted_request",
        optimizableInstruction = REDACT_INSTRUCTION,
    ) { _, input -> "user_query:\n${input.userQuery}" }
    
    val callUntrustedModel by subgraph(
        name = "untrusted_model",
        llmModel = untrustedModel,
    ) { /* one call to the model */ }

    val respondToQuery by optimizableSubgraphWithTask(
        name = "respond_to_query",
        optimizableInstruction = RESPOND_INSTRUCTION,
    ) { _, input -> input.rendered() }

    // …and the same plumbing and wiring as before.
}

With the step marked as “optimizable”, the graph itself remains unchanged. The instruction just moved into a parameter of its own, where an optimizer can find it and replace it. When there is no artifact to load, the step runs on the text you wrote.

Running the optimizer

GEPA runs your agent over the training set, shows a second model what went wrong, and asks it for a better instruction. Then it tries that instruction and repeats it, round after round, within the budget you set. Across those rounds it keeps several candidate instructions alive at once, so it does not drop one that only wins on your hardest items.

The optimizer needs one thing from the agent’s author: a function that, given a finished run, returns a score and a short summary of how it went. The score decides which candidates survive, and the text is what the second model reads. Here, that text is one line of numbers:

PUPA overall score: 0.50. Quality score: 1.00. Privacy score: 0.00. PII leakage ratio: 1.00.

That line, plus what the step itself received and produced, is everything GEPA had to work from when it wrote the instruction in the next section.

The rest is just two calls:

val optimizer = GEPAOptimizer(
    gepaFeedback = pupaFeedback,
    reflectionModel = OpenAIModels.Chat.GPT4_1Mini,
    numRollouts = 2426, // a practical start is 6 to 20x your training set
    storagePath = ResilientPath("build/artifacts/pupa-gepa.json"),
)

optimizer.train(session) // the session wraps your agent, dataset, and metric
val optimized = optimizer.loadOptimizedAgent(agent)

train writes the artifact, and loadOptimizedAgent reads it back and hands you a copy of your agent running the new instructions. Every optimizer in the library has those same two methods, so switching between them leaves the surrounding code alone.

What GEPA wrote

The step that strips the personal details started with this instruction:

Given a private user query, create a privacy-preserving request for a powerful external LLM. The LLM may assist without learning private information about the user.

Two lines, deliberately vague. It says nothing about how much redaction would be enough.

By the end of the run, that same step was running a 4,456-character instruction. Here are three pieces of it:

2. Redaction and Generalization
   - Replace all such information with broad, generic placeholders or neutral terms
     (e.g., "the minor," "the course," "the company," "a noted musician," "two large
     urban communities," "the agency," "the program," or "a fellowship honoring an
     individual").

Additional Domain-Specific Considerations:
- For academic or educational content, avoid exposing real professor names, specific
  course titles or codes, individual student names, or direct internal feedback. Use
  neutral terms like "the minor," "the course," "professors," "students," and "mentors."

Overall Strategy:
Create a balanced abstraction that ensures zero leakage of private or identifiable data
while preserving the user's original task in sufficient detail for accurate execution
by a LLM.

Every detail in the instruction above came out of the runs, not from us. You would probably not sit down and write the placeholder vocabulary, the rule about professor names and course codes, or the overall strategy at the end. That strategy states a trade-off we never mentioned: keep the leakage at zero, but leave enough detail to do the job.

Results

We wanted to know whether our GEPA in Koog does what it is supposed to, so we ran it on the benchmarks the project itself uses and checked whether the agents performed better.

Take PUPA, for example, the benchmark that the above agent runs on. Its baseline was already 81.19, with little room left above it, but optimization still improved it by 13 points to 94.22. Both numbers are the mean of five runs, and the gain is statistically significant.

The rest of the benchmarks, including how our numbers line up with the original implementation, can be found in the docs.

Three more optimizers

Besides GEPA, the library implements three more optimizers behind the same interface:

• BootstrapFewShot mines demonstrations out of your successful runs. It makes no LLM calls beyond running your own agent, making it the cheapest option to try first.
• MIPROv2 searches across instructions and demonstrations.
• ACE grows a playbook of insights that rides along with the system prompt, for when you want the agent to accumulate knowledge rather than a better-worded instruction.

Swapping between them is one constructor; the train and loadOptimizedAgent calls around it stay as they are. Writing your own means implementing those same two methods, and all four of ours fit behind them without an escape hatch.

Where to find it

The library, its documentation, and the full benchmark results are in the open-source repository. The optimizable-subgraph support it relies on currently lives in a fork of Koog; we expect it to land upstream.