Prompts That Can Write Themselves: A Kotlin Experiment

JetBrains Research released the Koog Agents Optimization Library, a research library that tunes the prompts of Koog agents without requiring the agent to be rewritten, and reported that the GEPA prompt optimizer raised its PAPILLON privacy agent's score on the PUPA benchmark from 81.19 to 94.22. Developers mark steps worth tuning with the optimizableSubgraphWithTask builder, supply a training set and a scoring method, and the library returns improved instructions as a separate JSON file that can be reviewed and committed. The library is published on GitHub with documentation and examples.

JetBrains Research Research is crucial for progress and innovation, which is why at JetBrains we are passionate about both scientific and market research Research https://blog.jetbrains.com/research/category/research/ 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 https://github.com/gepa-ai/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 https://huggingface.co/datasets/Columbia-NLP/PUPA . GEPA improved our agent’s score from 81.19 to 94.22. 👉 Koog Agent Optimization Library GitHub https://github.com/JetBrains/koog-auto-agent-optimization-lib 👉 Docs https://jetbrains.github.io/koog-auto-agent-optimization-lib/ 👉 Examples https://github.com/JetBrains/koog-auto-agent-optimization-lib/tree/main/koog-optimization-examples An agent to optimize Koog https://github.com/JetBrains/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 https://arxiv.org/abs/2410.17127 , a privacy pipeline that GEPA uses as one of its benchmarks. It runs on PUPA https://huggingface.co/datasets/Columbia-NLP/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<PupaAgentInput, PupaRedactedRequest name = "craft redacted request", { input - REDACT INSTRUCTION + "\n\nuser query:\n${input.userQuery}" } val callUntrustedModel by subgraph<String, String name = "untrusted model", llmModel = untrustedModel, { / one call to the big model / } val respondToQuery by subgraphWithTask<PupaResponseInput, PupaFinalResponse 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<PupaAgentInput, PupaRedactedRequest name = "craft redacted request", optimizableInstruction = REDACT INSTRUCTION, { , input - "user query:\n${input.userQuery}" } val callUntrustedModel by subgraph<String, String name = "untrusted model", llmModel = untrustedModel, { / one call to the model / } val respondToQuery by optimizableSubgraphWithTask<PupaResponseInput, PupaFinalResponse 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 https://huggingface.co/datasets/Columbia-NLP/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 https://jetbrains.github.io/koog-auto-agent-optimization-lib/ . Three more optimizers Besides GEPA, the library implements three more optimizers behind the same interface: - BootstrapFewShot https://dspy.ai/api/optimizers/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 https://dspy.ai/api/optimizers/MIPROv2/ searches across instructions and demonstrations. - ACE https://arxiv.org/abs/2510.04618 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 https://github.com/JetBrains/koog-auto-agent-optimization-lib . The optimizable-subgraph support it relies on currently lives in a fork of Koog; we expect it to land upstream. Subscribe to JetBrains Research blog updates