Training a Specialist Code Search Agent with turbopuffer

Today’s frontier models are strong code searchers armed with nothing more than the tools a human would use: the filesystem, ls, and grep. That is usually enough on a single small repository, which is why coding agents generally use search tools that don't rely on a precomputed search index. However, this approach to code search becomes expensive and slow when the corpus is large. In these cases, savings from eschewing search indexes are offset by high inference costs for frontier models and slow grep tool calls.

We partnered with turbopuffer to show how to solve both of these problems:

1. Train a small, specialized search model. Today’s small open-weight models struggle with search tasks out of the box, but they can be post-trained with search tools to achieve search quality near current frontier models at a fraction of the cost.
2. Build a dedicated search index in turbopuffer. On large corpora, search tools backed by a precomputed index return results several orders of magnitude faster than grep, and the gap only widens as repositories are added.

Training setup

We create an RL environment to train Qwen3.6-35B-A3B to use search tools to achieve different search tasks across a large multi-codebase corpus.

Data and task construction

We index ~9,000 real open-source GitHub repositories spanning many languages and project types. Each training task has a random subset of codebases assigned to it, including the target codebase(s) needed to complete its task. We test two task types to evaluate whether the agent can search both deeply and broadly:

Narrow task: the agent is given a codebase of up to 1,000 repositories and asked a question about one of them. The target repository is anonymized and the question is specific enough that arriving at the correct answer requires reading code and understanding implementation, not just skimming documentation, for example:

“We have a stateless database dump tool that can compress and encrypt dump parts before storing them. What is the exact order of compression/encryption on write and decryption/decompression on read, and how is the AES key material and nonce derived from the configured key string?”

The narrow task tests the agent’s ability to pinpoint the target repo in a very large set of codebases (a needle-in-a-haystack search) and then pull exact details out of its code.

The reward is a gated combination of:

• Correctness: each question we generate comes with a rubric graded by an LLM-as-judge (GPT-5).
• Citation accuracy: the model is asked to make specific citations to code blocks in its answer; we grade it for how well its claims are supported by these citations alone.
• Efficiency: we applied a penalty in the reward term to incentivize short answers reached in few agent turns (more below).

Open-ended task: the agent is given a corpus of up to 100 repositories and asked to find examples within the corpus which implement some described code pattern (usually a design pattern or architectural idiom). Only a handful of the repositories in the corpus actually have the pattern, and the task is to find them and confirm each one by reading its code, for example:

“Find as many examples of the ‘visitor/double-dispatch’ pattern in the corpus as you can. This is a code pattern where operations are externalized from a class hierarchy via double dispatch; each node exposes accept(visitor) and calls back the visitor's type-specific visit method, so new operations are added as visitors without modifying node classes.”

This task rewards the ability to efficiently sweep the whole corpus, and to search in concept space rather than code-implementation space. Matches are on an abstract shape and it is much harder to identify specific strings that an agent might grep for, which should make this task very challenging for a bash-only agent.

The reward for the open-ended task is precision-focused. Each repo cited in the answer is assessed independently by an LLM judge that sees the cited code in context, and counts it only if the citation resolves to real code and the judge agrees it implements the specified pattern. The task score is precision — the fraction of claimed repos that are correct — multiplied by the same efficiency penalty as the narrow task:

We build each open-ended task so that (at least) K of its repos implement the pattern, and the model is told that there are K repos to find. The expression in the denominator is chosen to encode both precision and recall: finding fewer than the K total caps the score while padding the response with junk repos hurts it. This ensures the agent must search broadly but is only incentivized to return repositories it feels sure about.

Both tasks apply the same efficiency penalty: it rewards short answers reached in few agent turns, and is the average of a length term and a turn term:

where $L$ is the answer length in characters and $n$ the number of agent turns. The length term gives full credit up to a length $L₀ = 2048$ and decays to $0$ by $L₀ + B (B = 6144)$; the turn term decays from 1 toward a floor $f=0.4$ as turns approach a budget $N=20$ for the narrow task and $N=40$ for the open-ended task. An empty answer scores 0, and one that never searched (n ≤ 1) is capped at 0.1.

Harness

The model was given a code search harness inside a read-only sandbox. The baseline harness provides a standard set of read-only exploration tools to navigate the filesystem: read_file (reads a file, optionally a specific line range), ripgrep (recursive regex search), glob (matches files by glob pattern), and list_dir (lists a directory’s contents). We then additionally provide tools to perform hybrid search over precomputed indexes stored in turbopuffer. We use tree-sitter to chunk repositories in an AST-aware fashion, then index each chunk in two ways:

1. a BM25 keyword index
2. a dense embedding of each chunk, produced by the open-source Octen-8B embedding model

Alongside the text and vectors, we store the filepath and line numbers of the chunk as metadata. At search time, we retrieve candidates with turbopuffer, rerank them with the open-source Qwen3 Reranker 4B, and use the metadata so the agent can jump straight to the source file and cite it directly. The index itself is query-agnostic: HyDE (below) changes what we embed at query time, not how the index is built.

Search tool modes

We train and evaluate our agent under two search modes:

1. bash-only: the agent has access only to the filesystem and bash primitives
2. bash + turbopuffer search: in addition to bash primitives, the agent can search the precomputed turbopuffer indexes. Rather than embedding the question directly, the agent writes a hypothetical code snippet — its guess at what the answer looks like in source — which we embed with the document prompt, fuse with a BM25 leg over the same snippet, and rerank. This is HyDE⌝ applied to code search: a query written in code sits far closer to the code being searched than a question in English does.

Training Results

We blend the two task types in equal proportion and train Qwen3.6-35B-A3B with GRPO.

Training increases search quality

Regardless of task or tools, the model becomes more effective at completing the search task as it learns to use the tools available to it. On the narrow task, the bash-only agent improves its correctness × citation support by 38% over training, while the agent with turbopuffer search improves by 57%. On the open-ended task, the bash-only agent increases precision by 182% and the turbopuffer agent by 211%. Training helps either way — but the agent with access to turbopuffer search tools is learning from better evidence at every step, and it finishes 16% ahead on the narrow task and 25% ahead on the open-ended one.

Turns decrease over training

In addition to achieving better search results through training, the agent also learns to do so more efficiently, reducing turns over the rollout. This is especially pronounced in the open-ended task, where the bash-only agent ends training using 50% fewer turns than it did at the start, and the agent with turbopuffer search tools uses 70% fewer.

Tool use shifts toward search

Reducing turn counts show that the agent gets more efficient over training. We can see exactly how this occurs by looking at the mix of tool calls as we train:

Training pushes the two agents in opposite directions. The bash-only agent responds by grepping harder — on the narrow task its ripgrep calls more than double over training (4.0 → 8.9 per rollout) and its total tool calls actually rise, from 11 to 12. The agent with turbopuffer search goes the other way: ripgrep all but disappears (1.7 → 0.2 per rollout on the narrow task, 15.3 → 0.3 on the open-ended one) and total calls are reduced by half on the narrow task and by 70% on the open-ended one. It also learns to search better rather than more — turbopuffer search calls on the open-ended task drop from 12 per rollout to 5 even as precision climbs.

Usage of filesystem primitives like glob and list_dir decreases in both modes, but only the search agent can afford to give up ripgrep: for the bash-only agent ripgrep is the effective search index, and it finishes training sending 70% of its calls there.

Both agents also learn to fire several tool calls in a single turn. By the end of training an open-ended rollout averages 2.6 (turbopuffer search) to 3.1 (bash-only) tool calls per turn, which is why turn counts fall faster than call counts do.

Ranking each rollout’s calls from first to last at the final checkpoint shows two different strategies:

• The bash-only rollout opens with ripgrep — 97% of first calls on the narrow task, 99% on the open-ended one — and only turns to read_file near the end.
• The turbopuffer rollout opens with a search instead (96% narrow, 77% open-ended) and spends the rest of the rollout reading: by the last position decile, 88% of its narrow-task calls are read_file.
• On the open-ended task the search agent keeps searching throughout, with a quarter of its final-decile calls still going to turbopuffer search. Sweeping a corpus for a pattern means alternating search and read, rather than locating a single target and settling in to read it.

Comparison to frontier models

We compare the base Qwen3.6-35B-A3B checkpoint to its trained counterpart as well as several frontier models to determine the pareto frontier of quality versus latency and cost. Before RL, the Qwen3.6-35B-A3B is not a very good search agent. The training is what makes the specialist.

Our fine-tuned model tops the narrow task outright in correctness × citation support. While it performs less admirably in the open-ended task, it still closes the gap considerably versus the base model. In both cases the specialist runs at 10-100x lower cost and 2-10x lower latency than most frontier models.

Latency

One of the most compelling reasons to employ search tools over precomputed indexes is how efficiently queries scale over the size of the corpus. This holds true in our training setup. Below we show the latency distribution on the open-ended task across different search-tool modes, when we scale the size of the corpus from 20 repositories to 300 repositories.

We observe that model completion-time is roughly similar regardless of the size of the repository or the tools available. In contrast, a 15x increase in corpus size causes ripgrep latency to increase by 11x, while turbopuffer search latency increases by only 1.2x.

Cost

We showed that a small model trained to use search tools over precomputed indexes can run search tasks at up to 100x lower token cost than frontier models. It is of course important to note that the cost of the full search pipeline will be determined by several factors beyond inference:

• The size of the corpus. Larger corpora are generally more expensive to store, index, and query.
• The write-read ratio. Precomputed indexes effectively cache, at write time, compute that would otherwise be spent at read time. When the corpus changes infrequently and sustains a high rate of queries, the cost of indexing is well-amortized across many fast and cheap queries. When that ratio flips, however, the cost of maintaining the index can outweigh the savings at query-time.
• The kind of search tools used. BM25 full-text search requires little upfront compute, and zero inference, whereas dense vectors require compute for embeddings and also result in storage amplification. Use of techniques such as late interaction or hypothetical document embeddings (HyDE) can further increase inference and storage costs.

Conclusion and demo

Filesystem primitives are effective for code search agents using frontier models in a single-repository setting. The arithmetic changes for search at scale: once a workload involves many searches over large volumes of code, bash-only search can become a latency and cost bottleneck. Our results suggest a solution to this problem: train a small model to search an index. At large scale that combination makes the median search nearly 3x faster and cuts the marginal cost of a search by 100x.

Try it yourself in our demo⌝.