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

Part 1: Parsing, chunking, and vectorization

Some time ago, we set out to build the best semantic code search platform we could: a RAG pipeline that gives LLM agents precise, citable evidence from real repositories instead of whatever grep happens to surface. The eventual solution was JetBrains Context. We got it working, we got it into production, and we collected a lot of scar tissue along the way. In this series of posts, we’ll share the parts we wish someone had told us on day one.

Coding agents are undoubtedly the biggest technology leap for software development of our decade. Agents and frontier models are proving their aptitude in the face of seemingly insurmountable code complexity to produce ostensibly reliable code.

However, as more and more development processes become agent-driven, the agent’s efficiency and the quality of the produced code become increasingly important. The question is not so much about whether an agent can complete the task, as given enough time and token resources, it surely will, but rather how much time, effort, and steering is required for it to generate production-grade results. For large-scale code bases specifically, the agent would spend a great deal of time searching for the relevant pieces of code relevant for the feature it’s working on and pulling them into the context.

Why semantic search matters

Attempting to locate the right code snippets, the agent will resort to traditional tools for code search such as keyword search and grep. These tools, however, are limited in that they require the agent to know in advance which exact text to search for. For example, an agent looking for where session tokens get refreshed cannot rely on the code helpfully containing the word “refresh”. To reason through abstract domains, the agent needs the ability to search for code by meaning, also known as semantic search. This is where retrieval-augmented generation (RAG) comes into the picture. If we can index the source code in a way that captures its semantics and then allow the agent to retrieve the relevant pieces on demand using free text search, we create an interface that plays to the agent’s strengths.

From prototype to production

Like many great ideas in the agentic era, a native, prototype implementation is extremely simple. A well-evaluated production grade solution most certainly is not. In this series of blog posts, we want to share what is involved in making an effective RAG system, as well as the wrong turns we took in our journey to create our own: JetBrains Context. We’ll tackle each stage, from pre-processing to storage and agent integration, providing some more technical context and advice.

This first part of the series will cover the initial stages of the pipeline: parsing and chunking, where raw source files are divided into properly scoped units, and vectorization, where those units are transformed into a representation that supports semantic search.

The fine AST of parsing and chunking

Parsing and chunking is a critical pre-processing step in a good RAG solution, but it is often overlooked. In order to allow the LLM to embed or otherwise index the source code, we must first feed it the raw lines of code. This may sound trivial, and probably would be for small-scale demo projects. However, production-grade systems contain thousands of files, which, in turn, span hundreds or even thousands of lines. If anything, agents have compounded the problem, as they tend to be prolific writers, further inflating the codebase. Each file may contain multitudes of classes, fields, and methods, with varying degrees of relatedness among them.

Finding the right chunk size

Even if it were possible to fit these huge code files into an embedding model in their entirety, that expensive feat would ultimately be self-defeating. Because the entire file was embedded in a single unit, the search would return the entire file. This is counterproductive to the goals of agentic code exploration and navigation, which are mostly concerned with finding a specific function, symbol, or code snippet.

On the other hand, if we were to take the other extreme and granularly embed each separate line of code, we would be facing a problem of a different sort. These individual lines can be semantically insignificant without the surrounding context. A generic function name or comment does not merit embedding and will produce the wrong retrieval result. In a sense, we would not be able to see the forest for the trees, and the agent would be overloaded with multiple, often insignificant micro-results.

It is therefore imperative to find the right method to chunk or divide the code into groups that are properly scoped. Each group should include enough of the necessary context and represent common semantic meaning.

Why fixed-size chunking falls short

Chunking is a generic name for the technique of taking content that will be fed to the agent and dividing it into a set of chunks. A naive approach to chunking could be simply splitting a large file into groups with a fixed number of lines. However, if we were to take that approach, we would find the resulting groupings semantically wrong. Unrelated code pieces would be grouped together, for example, an import statement and some function content, leading to mistakes during retrieval.

To solve the problem, we can leverage the fact that every source file has a pretty well-defined structure. Take Java as an example – imports tend to be at the top of the file, followed by a class definition with an optional doc-comment preceding the header. The class will contain fields and methods, which in turn may also have their own doc-comments. Knowing about the conventions and rules that define the class structure allows us to perform smarter chunking and achieve the right balance of surrounding information.

Parsing and structure-aware chunking

Over the last 26 years, we at JetBrains have developed parsers that are smart enough to adjust for the various quirks, irregularities, conventions, and nuances of specific languages. Alongside other tools, these parsers form our internal JetBrains Code Engine platform on which JetBrains Context is developed. At the moment of this article’s composition, JetBrains Context supports parsing and structure-aware chunking for nine major languages: Kotlin, Java, Python, JavaScript, TypeScript, C#, PHP, Go, and Rust. For all other languages, our implementation simply falls back to naive, line-based splitting to ensure that any language or document can be indexed and searched.

The parser allows us to break source files into streams of syntax nodes that carry information about what they represent – comments, whitespaces, lists of modifiers, and so on. The chunking algorithm then consumes that stream and applies logic that decides the scope of a given chunk. Based on the node’s type and size, as well as its descendants, the algorithm makes a decision. If a node exceeds the size threshold but has no children, it will fall back to more primitive splitting strategies.

Some language-specific constructs are kept as single slices even if they exceed the preferred size. Prefixes such as documentation, annotations, visibility modifiers, and keywords are kept together with the declaration; suffixes (usually closing syntax) remain associated with the construct they close. There is also some language-specific cleaning, where, for instance, common and semantically meaningless Java annotations such as @NotNull or @Override are removed.

The algorithm bears some similarities to cAST, authored by Zhang et al. in 2025. Both our implementation and cAST retain the largest syntax units that fit, subdividing only the units that are too large, and grouping smaller adjacent units to avoid tiny chunks that are not usually semantically meaningful. The biggest difference is that we coded more language semantics into our implementation, keeping Python decorators together with definitions, KDocs next to Kotlin declarations, and so on.

After grouping, chunk normalization is performed, which involves:

Trimming leading and trailing whitespaces
Deleting blank lines
Removing common indentation while preserving relative indentation

Following the normalization procedure, the chunk is then passed to the next step – embedding – along with metadata that consists of a relative path, which gets embedded alongside the normalized chunk content.

Evaluating the quality of chunks

It is hard to give a concrete answer as to what the input to the embedding model should look like. Chunk size matters, but as discussed before, bigger is not always better. Additionally, some metadata embedded alongside the code may be useful, while some may introduce noise that ultimately decreases search quality.

We opted to use an LLM-as-a-judge strategy to inspect the chunks as a part of the evaluation. The judge, using a chunk and the source file, considers whether the boundary makes sense. It looks for unexpected artifacts, such as detached documentation, orphaned closing syntax, or fragments of code that are cut through a meaningful construct. In addition, any changes to the source code processing pipelines also go through the full, end-to-end retrieval evaluation. We’ll get back to that evaluation pipeline in the following part of this series.

Vectorization

Having pre-processed the source code, we finally have text chunks that are hopefully just the right size and correctly grouped for semantic retrieval. Our next task is to transform these fragments in a way that will later allow us to support semantic search, through a process called vectorization.

With vectorization, an embedding model reads a piece of text and emits a fixed-length list of numbers (a vector), which amounts to a point in a space of a few thousand dimensions. Significantly, the model is trained so that texts with similar meaning land close together. Traditional search might miss the connection, but here, a function that flushes buffered write operations and one that drains a pending queue can end up near each other despite sharing no common keywords. The distance between vectors hence becomes a measure of relatedness. A query is turned into a position in the same space, and the results are whatever lies nearest to it.

Punch for the byte: Optimizing for storage

Any attempt to vectorize a large codebase must take into account both cost and performance. A single embedding is cheap, but a large repository produces millions of chunks, which become millions of vectors that must be stored, held in memory, and compared against each incoming query. A vector of a few thousand dimensions in 32-bit floats weighs around 16 kilobytes, so a few million chunks add up to tens of gigabytes of index before any bookkeeping. At such a scale, the allocation of bytes per vector becomes cost-limited, and the leading question quickly shifts from “how accurate can we be?” to “what do we get per byte?” In other words, we need to find a way to reduce the cost while retaining as much search quality as possible.

There are two ways to reduce vector cost. The first is to keep fewer dimensions. Modern embedding models are trained so that a leading slice of the vector works on its own. The dimension loss is applied across several nested prefix lengths simultaneously, pushing the coarsest structure into the earliest dimensions. This means you can cut a vector short and renormalize it, and it still retrieves. Alternatively, you can keep every dimension and spend less on each one by sacrificing on precision and thus keeping fewer bytes for each vector.

These two options are independent of each other and can be combined, which means any storage budget can be met through different mixes of dimension count and numeric precision. The real question is which mix retrieves best for the same number of bytes. The trade-off is far from even. Suppose the budget is 512 bytes per vector. You could spend it on 128 dimensions kept at full 32-bit precision, or on all 4,096 dimensions kept at a single bit each. Both fit the budget exactly, but in testing, you’ll find that the second option retrieves considerably better.

Why dimensions matter more than precision

To see why, it helps to think of each dimension as one small question the model has learned to ask about the text: Is this about error handling? Does it touch the network? Is it test code? And there are a few thousand similar topics and questions that haven’t been named. (The real dimensions are blurrier than that, but this is a useful abstraction.)

No single answer means much on its own. We consider two chunks to be similar when their answers to many of these questions are the same. Therefore, we should assess the vectors by looking at the coverage of the questions rather than the exactness of the answers.

Keeping all 4,096 dimensions at one bit preserves a rough yes-or-no answer to every question. Truncating to 128 dimensions keeps very precise answers to three percent of the questions and throws the rest away, and no amount of precision on the surviving dimensions can recover the information the discarded ones carried. In a sense, a long questionnaire filled in with checkmarks beats a short one filled in to six decimal places. Dimensions are what you want to keep; precision is what you can afford to lose and is easier to compensate for later on.

So we chose to keep every dimension and take the precision reduction to its limit, dropping the vectors to one bit each, which is 32 times smaller than the same vector in 32-bit floats. The quantization itself turns out to be surprisingly simple. Every component at or above zero becomes a one, while every negative component becomes a zero, and the magnitudes are thrown away:

Changing the representation changes the metric with it. Cosine similarity needs the magnitudes we just threw away, so binary vectors are compared by Hamming distance instead, which is simply the number of positions where two bit patterns disagree. Compare, for example, 10110100 and 10010110. They differ in two positions, so the distance between them is two. At full length, the computation stays just as simple. A 4,096-bit vector is stored as 64 words of 64 bits, and comparing two of them means XORing each pair of words, which leaves a 1 wherever the two vectors disagree, and then counting the 1s. A CPU does each of those in a single instruction per word, so a full comparison costs in the order of a hundred instructions where cosine similarity on the original floats needed thousands of multiplications.

Note that the metric was never a separate decision. We chose one-bit precision for the storage savings, and once every component is a sign bit, Hamming is the only comparison left that makes sense. Choosing the precision chose the metric.

Binary quantization still costs a few points of recall against the unquantized vector. We accepted that cost after considering that a reasoning agent would be consuming the results. A code search feeding an agent needs the right neighborhood far more than a perfectly ordered top 10. When the agent asks where session tokens get refreshed, what matters is that the relevant handful of files shows up among the first dozen results. Whether the best chunk ranks second or fifth changes nothing, because the agent opens the candidates and reads them anyway. In that loop, a ranking degradation that would be plainly visible in a three-result UI built for humans is mostly invisible.

The limits of binary quantization

The trade-off we made had a subtler cost that took us a bit longer to understand. Binary quantization doesn’t only sacrifice accuracy; it compresses the *range* of similarity scores. With full-precision vectors, an unrelated pair can score near zero while near-duplicates score near one, a comfortably wide spread. Sign bits behave differently. Around half the bits of two entirely unrelated vectors still agree by pure chance, while a strongly related pair might have agreement for two-thirds. So every score in the index, relevant or not, lands in that thin band.

Ranking survives the compression, since relevant results still score above irrelevant ones, but thresholding does not. Picture a feature that volunteers related code without being asked, say a panel that suggests existing implementations while you type. Its most difficult requirement is knowing when to stay silent. To make that determination, it needs a usable gap between “related” and “unrelated” scores. Binary vectors don’t leave one. Any cutoff placed inside that narrow band either fires on everything or on nothing. So where an index needs an absolute relevance judgement rather than a relative ordering, we keep 16-bit floats and pay for the storage.

Embedding scope

While indexing and searching use the same model, the two jobs could not be more different. Indexing is throughput-constrained, with millions of chunks asynchronously handled. The GPU will handle about 32 chunks per batch before becoming saturated. A search, on the other hand, needs to be fast and responsive. Users will give up if they are not provided with results within a couple of seconds at most. Therefore in deploying these models we optimize them accordingly: one to maximize chunks per second, the other for minimizing time to first result.

We chose an instruction-following model, trained with a deliberate asymmetry between the two sides of retrieval. Significantly, the two sides are represented by very different types of text. A query is a short question in natural language, while a document is a chunk of code. A document is embedded as is at indexing time. A query is wrapped with an instruction describing the retrieval task, something like “given this search query, find the code that answers it”, which tells the model what role the text is playing. We preserve that arrangement at inference because it is the shape the model learned.

To allow the two sides to align more easily, we embed each chunk together with its file path. The path supplies metadata that the chunk alone lacks: which module it lives in, and what the file is. In a monorepo, though, the path itself becomes a problem. The IntelliJ IDEA monorepo runs to over a million files. The median source file there sits nine directories deep behind a 91-character path, and close to 10,000 source files have paths longer than 150 characters, the longest of them 218. That is before any checkout root is prepended.

Most of those characters are used for structural nesting and offer no useful information about the file. A run of segments like `src/org/jetbrains/kotlin/idea/k2` restates the package hierarchy, which a compiler needs and a search does not. Meanwhile, the file at the end of that longest path is 24 lines long. If we simply embed the path text as is beside a chunk, we’ll find that the path will sometimes take up more space than the code itself. To compensate for that, a path is capped before it reaches the model, and the rule is that *both ends survive*. The leading segments tell you which module you’re in, while the last two, the immediate parent and the filename, tell you what the file is. The middle is the part that can go, and only as much of it as the cap requires. Keep the longest prefix that still fits, elide what falls between into `…`, and if even parent-plus-filename is too long, keep only the name itself.

The same discipline applies when a user scopes a search to a subdirectory. The obvious implementation is a metadata filter: run the search as usual and discard results that fall outside the directory. We do something different. The scope is rendered into the query text itself, in the same shape, with the same abbreviation function and the same separator the indexed chunks used. If a chunk went into the index under the abbreviated form of `community/plugins/kotlin`, a query scoped to that directory carries the same string in exactly the same form, so the query vector lands in the same region as the chunks it is supposed to match.

Protecting source code

There was one last design consideration we took into account. It was important for us to be attentive to customer privacy and security concerns. The source code of a company is often the core of its IP. Exposing it to third-party cloud models, or even to another company, increases the risk of inadvertently exposing sensitive data or even training other models to use it.

To make sure we address these concerns, we made the decision to adhere to several practices early on:

Avoid storing the code in our systems: A chunk holds a cluster reference, an item type, a file path, start and end offsets, a reference to a vector, and an optional metadata field. No content, no copy of the source code itself, is saved. What a search returns is coordinates, and the snippet you see is assembled on your machine, from your checkout, using them. The server just knows that something relevant lives at bytes 4,102–4,890 of a given path, not what it is.
Don’t use data for training: Every code index JetBrains Context builds is embedded by an open-weight embedding model, running on GPUs we operate. No embedding request leaves our infrastructure – not to OpenAI, not to Google, not to any other vendor. Therefore, we can guarantee that none of the data will be used to train anything.

These self-imposed design restrictions carry no cost in terms of retrieval quality. We evaluated the open-weight candidates against the hosted embedding APIs from the major providers on our own code-retrieval benchmarks, and ours came out on top. Open-weight embedders are now good enough that the interesting engineering has moved into what you feed them, how you serve them, and what you choose to keep.

A summary that is an interlude

In this blog post, we covered the first stages of the retrieval pipeline: the journey from raw source files to compact vectors that are ready to be searched.

At this point, we have millions of binary vectors and a way to produce more. The problems we haven’t solved yet are how to store them efficiently, how to create a system that can answer a query in milliseconds, how we can continuously evaluate our results to ensure we are making the right choices, and how we can get the agent to actually use our shiny RAG apparatus.

These topics and more will be the subjects of the next parts in this series, which we’ll be releasing over the next few weeks. As always, please feel free to ask any questions in the comments or share your own hard lessons from designing a RAG solution. We are eager to learn of different and creative ways you have found to be effective! In the meantime, feel free to check out JetBrains Context, currently in public preview, it is already included with your JetBrains license 😀

Until next time!

Discover more

Fuente: Artículo original