Building a RAG Pipeline for Semantic Code Search: A Developer Diary and Field Notes

JetBrains detailed in the second post of its developer diary series how its production-grade Air Context RAG system for semantic code search stores and retrieves millions of code vectors, combining approximate nearest-neighbor indexes such as HNSW and IVF with a follow-up exact search pass over a narrowed candidate set. JetBrains said the hybrid approach "works remarkably well," letting the approximate index handle scale while the exact pass handles the last mile of precision. The post also covers filtered ANN, in which deterministic metadata filtering by repository, branch, path and user permissions must supersede the search algorithm, and incremental index updates as commits land without duplicating chunks across branches.

JetBrains AI Supercharge your tools with AI-powered features inside many JetBrains products Agentic AI https://blog.jetbrains.com/ai/category/agentic-ai/ AI https://blog.jetbrains.com/ai/category/ai/ Building a RAG Pipeline for Semantic Code Search: A Developer Diary and Field Notes Part 2: Storage, Search, and Agent Tooling In this blog series, we aim to share some of the lessons we learned over the past six months as we built and rolled out Air Context https://www.jetbrains.com/air/context , a production-grade RAG system for code search. In the first post of the series https://blog.jetbrains.com/ai/2026/09/building-a-rag-pipeline-for-semantic-code-search-a-developer-diary-and-field-notes , we covered chunking, parsing, and vectorization. Turning code into vectors, however, represents only part of the journey. In this post, we’ll tackle the next problems that must be solved for any serious production-grade RAG system: how to store millions of vectors so that the index keeps up with Git, and how to retrieve the right vectors in milliseconds. Storage Once we have the code parsed, chunked, and turned into vectors, our next question is where to put them. A real-world codebase can produce millions of vectors. To compound the problem, codebases are not static. Not only do they change often and exist in dozens of branches at once, but the agent expects an answer in milliseconds, filtered down to the right repository, the right path, and only the files that the user is allowed to see. All this requires a special storage layer that efficiently handles the search operations under the given constraints. At the center of this storage layer sits the vector index. Vector search requires finding the vectors closest to the one representing the query. Comparing a query against a million vectors one by one is far too slow, so in practice we use approximate nearest-neighbor ANN structures such as HNSW or IVF. These skip most of the comparisons and return results that are very close to the best match, but are not guaranteed to be the best match. That speed comes with trade-offs when it comes to how many of the truly relevant results are actually returned also known as recall , how much memory the index needs, how well it copes with a constant stream of inserts and deletes, and whether it can respect filters at all. On top of that, any healthy codebase must support multiple incremental index updates as commits land. To ensure this can happen at scale, the design must avoid keeping a separate copy of every chunk for every branch that shares the same code, and it must maintain consistency for the vectors and the metadata that describes them. None of these capabilities are provided by embedding models. Without strong storage solutions, our painstakingly extracted vectors will simply not be usable. The role of approximation Approximate search algorithms sacrifice some measure of recall, but they operate well over a large dataset. We found that a hybrid approach, in which we also employ exact search over a more limited vector space following the initial ANN search, works remarkably well. In other words, using a cheap first pass narrows the whole corpus down to a small set of likely candidates. Running an exact comparison then becomes affordable. The approximate index handles the scale, and the exact pass handles the last mile of precision. This gives us another lever by which to fine-tune the search, as we are thinking not only about which algorithms to use but how to use them in tandem. Filtered ANN: The nuances Vector search for semantic similarity is blind to harder search constraints, such as the limitation of results to a repository, branch, or specific subfolder. Additionally, security and permission constraints, such as restrictions on the kind of data that can be queried, must be strictly enforced. This type of metadata filtering needs a deterministic mechanism that will supersede any search algorithm. Combining a filter with an ANN index is genuinely hard because the two mechanisms fight each other. An HNSW graph is built assuming you can traverse freely between neighbors; the moment you constrain traversal to nodes that satisfy a predicate, the graph can become disconnected and the walk can miss valid candidates. There are three broad strategies to address this: 1. Pre-filtering. This means applying the metadata predicate first, and then searching only the surviving set. It provides 100% recall relative to the filter. However, it loses the ANN index over the filtered subset and effectively falls back to exact search. This strategy is best suited for highly specific filters, such as a single small repo or a narrow subfolder. However, it becomes expensive when the filter still leaves a huge population to search. 2. Metadata bundled with the vector index. The strongest general-purpose approach is a storage that makes filtering a first-class citizen of the index itself. In other words, parameters such as folder permissions or repos are searchable and affect the vector distance. Qdrant and turbopuffer are built this way. You keep most of ANN’s speed and honor the filter without falling back to brute force search. For a production codebase RAG with real ACL requirements, this would be a balanced choice. 3. Post-filtering. Run the ANN search first, then throw away results that fail the predicate. The problem, of course, is that because the top-K candidates were not selected with this filter in mind, most or all of them could end up being invalid, leaving you with a short, low-quality list. pgvector’s iterative-scan mode mitigates this by fetching more and more candidates until enough survive the filter, but it carries a cost in latency and can be unpredictable. Post-filtering is considered to be the weakest of the three strategies. Keeping up with incremental code changes When working on large code repositories, a RAG system can’t afford to re-index the entirety of the codebase upon every change. Thus, to achieve viable indexing performance and data storage limits, indexing must be performed incrementally. Ideally, we’d want each repository snapshot to be deterministically clustered into a Merkle tree mirroring its directory structure. New revisions are diffed from the top down, and entire unchanged subtrees are linked to the new snapshot as-is, without any reindexing. The embeddings are further deduplicated by chunk content hash, so within the modified clusters themselves, we would only index the changed code. With incrementality in place, index updates and deletions run step by step in small, lock-bounded transactions with limited parallelism, and heavy maintenance such as vector-index builds is offloaded to background jobs, preventing database CPU bottlenecks. However, all of this processing and business logic costs time and processing resources. Constructing HNSW vector indexes is one of the primary CPU bottlenecks in our database pipeline. To prevent system overload, we decouple vector ingestion from index building. Vectors are written first, and the HNSW index is built only once a repository’s snapshot is fully stored in the database. The build then runs asynchronously under strictly controlled parallelism. To keep CPU usage efficient, we use non-blocking concurrent index creation to allow at most one index build at a time per shard. This keeps the database instances stable, as a snapshot only becomes searchable once its index is ready. To deal with the CPU, memory, and metadata-capacity limits of a single database instance, our architecture scales horizontally across multiple shards, each backed by its own primary instance and read-only replicas. Search queries are served from the read-only replicas, deployed across geographically distributed regions, to guarantee low latency. Furthermore, to uphold data isolation, each customer is pinned to exactly one shard, with a simple application-side routing layer directing each transaction to the right shard based on the customer identifier. Leveraging local indexes for local changes In a typical SDLC, code is first written on a local laptop or desktop machine, later to be pushed to a remote branch and merged with the rest of the codebase. During these first development stages, that code is highly relevant to the individual developer, but less so to other developers working on their own branches. Creating a local index on the developer’s machine to support the pre-push development stage can help reduce the load and disk requirements from the central search index and ensure the search results are relevant. Air Context indexes can live fully on-device to support this use case. For local vector storage and search, we adopted USearch https://github.com/unum-cloud/usearch , a compact, high-performance engine built for vectors. Vectors and metadata are separated, with snapshots, clusters, and chunks kept in a dedicated metadata storage alongside the vector index. Because that metadata follows exactly the same model as the server, the entire incremental pipeline keeps working locally. Unchanged clusters are relinked across revisions, and identical chunks share a single stored vector. Sharing one metadata model between the server and local storage allows us to support cascade indexing, which permits the transition of code from local to shared. Search becomes a two-stage process. First, it runs against the remote index, and then it runs against the local one, using a reranker that merges the top candidates from both sources into a single result list. In this manner, we keep a flexible local index that is always synchronized with the local file tree, and at the same time optimize storage and inference by sharing remote indexes for remote branches. Furthermore, we save valuable time by transparently transitioning the index from local to remote as it gets pushed into the central source code. Post-retrieval optimization Post-processing: Turning candidates into an answer The nearest-neighbor search described above does not produce the final result. It produces a list of candidates: the chunks whose vectors sit closest to the query. While this is a great start containing a lot of useful data, it is not the final product to be delivered to the agent or user. The ANN results may contain near-duplicates, several fragments of the same function, or chunks that scored well for the wrong reason, for example because they share a few rare words with the query. Post-processing encompasses everything that takes place to transform the raw candidate list into the text the agent reads. In our pipeline, this covers removing duplicates, grouping fragments that belong to the same file, deciding how many results are worth showing, and, most importantly, reranking. Each of these steps is cheap compared to the search itself, because each one runs over a few dozen candidates rather than the whole codebase. Reranking Reranking can best be described as taking a second, more careful look at the candidates before handing them off. Our first pass has to be fast, since it compares each chunk to the query through two vectors computed independently. A reranker can afford to be slower. Its role is to look at the query and the candidate side by side and judge how well they actually match. While we can’t use an LLM to judge the relevance of results over a massive index, that does become a valid option when dealing with a more limited result set. Ultimately, reranking gives us a reordered list in which the best matches have moved to the top. In addition to this approach that uses neighbor search followed by reranking, we experimented with several other retrieval configurations. Some of them delivered a noticeable improvement in search quality, but at the price of significantly higher resource use and slower searches. One example is JointRank, a method presented at the ICTIR ’25 conference, which asks a language model to rank many candidates in a single pass by splitting them into overlapping groups, ranking each group in parallel, and combining the partial orderings into one. Agentic search Further attempts to improve search quality led us to explore using an agentic search tool. Instead of a single-shot retrieval, we can choose to employ a dedicated code-retrieval agent. This agent can orchestrate calls to several search and read tools, including the embedding-based retrieval. It can make judgment calls on which tools to call and use different thresholds based on the query itself. The agent decides on its own when to search, when to open a file, and when it has seen enough. We cap the number of steps it can take so that exploration never runs away with the budget. This retrieval agent can be implemented as a sub-agent, a common pattern established in recent years for helper agents that keep an isolated context. When a sub-agent does the reading, only its conclusion comes back, so the main agent’s context stays small and focused. Claude Code, for example, lets users define custom sub-agents that the main agent can call like any other tool. A sub-agent can also run on a cheaper or faster model than the one writing the code and can be restricted to read-only tools to prevent it from changing anything by accident. Additionally, several of them can run at the same time on independent questions. For our search tool, this means the main agent spends fewer tokens and less time on finding code and more on the task it was given. Result formatting and front-loading snippets A final aspect of the post-processing stage we have to consider is the format we use to hand the results over to the agent. Our goal is to save the agent time by preventing it from making additional calls, while also ensuring it has enough data to assess the relevance of the results for its original goal. This needs to be done judiciously, though, as the more data we front-load into the response, the more we inflate the agent context and possibly pollute it. The strategy we have adopted is to match the level of details to the assessed relevance of each result. The results for candidates whose relevance scores exceed a predefined threshold pack more information, including the code chunk itself as well as line numbers and additional metadata. For candidates with a lower estimated relevance, we’ll provide only some basic information like the file name and possibly a very partial snippet. This keeps the agent’s attention on the matches that are likely to be right while still hinting at where else it could look. From theory to practice In this post, we discussed the second part of the pipeline, covering storage along with search and retrieval. On the storage side, an approximate index does the heavy lifting, with an exact pass over a small candidate set. Filters for repository, branch, and permissions live inside the index rather than around it. A Merkle-tree view of each snapshot lets the index keep up with Git by reusing everything that has not changed, and a local index on the developer’s machine picks up the work before code is pushed. On the retrieval side, a reranker takes a second, more careful look at the candidates to judge their relevance. A dedicated sub-agent is capable of a more nuanced exploration when a single-shot search is not enough. Finally, the result format can provide the agent with details in proportion to how confident we are in each match, so the agent’s context stays small and focused. There’s still one critical line of questioning we haven’t yet addressed: Does any of this actually help the agent? For all of the optimizations, algorithms, and consideration we give to scale, does this RAG apparatus actually make a measurable impact on the agent’s performance? After all, agents, especially when using frontier models, have become increasingly adept at navigating large codebases. With that in mind, what is the added value of semantic search tools? These are important questions that have been on our minds from the very beginning of our journey. In the next post, we’ll show how we plug Air Context into agents, what it took to turn it into a product, and how we evaluate it – from retrieval benchmarks to end-to-end agent runs. Spoiler: in those runs, agents with Air Context needed up to 70% fewer turns to solve a task. As always, ask questions in the comments or share your own hard-won lessons. If you’d like to try it on your own code, Air Context https://www.jetbrains.com/context/ is in public preview and already included with your JetBrains AI subscription. Until next time Subscribe to JetBrains AI Blog updates