Building a RAG pipeline for semantic code search
JetBrains details the architecture of Air Context, their production RAG pipeline for semantic code search, focusing on AST-aware chunking and 1-bit binary quantization.
read it A detailed engineering deep-dive into production RAG trade-offs, featuring concrete architectural decisions like AST-aware chunking and 1-bit binary quantization.
summary
JetBrains shares engineering insights from building Air Context, their production RAG pipeline for semantic code search in coding agents. The post details how they moved beyond naive line-based chunking by using language-specific AST parsers to properly scope code units. It also explains their optimization of vector storage using 1-bit binary quantization, trading precision for dimension coverage and fast Hamming distance comparisons. Finally, it highlights their privacy-first approach of keeping source code off servers while outperforming hosted APIs with local open-weight models.
key points
- Language-specific AST parsers handle chunking for nine major languages, keeping declarations, doc-comments, and decorators grouped while splitting oversized units.
- Binary quantization (1-bit precision) reduces vector storage by 32x, enabling fast Hamming distance calculations via CPU bitwise instructions at a minimal cost to agent retrieval recall.
- File paths are aggressively truncated to keep only module prefixes and filenames, preventing deep monorepo structures from dominating embedding tokens.
- Architecture avoids storing raw source code on central servers, relying purely on coordinate references resolved locally.
who should read it
Software engineers and ML practitioners building retrieval systems or working with LLM coding agents.
discussion mixed
I think something like this would be key to improving the quality of coding agents. A very common issue (maybe the biggest one) observed by many people is that agents often produce a lot of duplicate and redundant code; multiple abstractions, methods, classes, data structures, etc. serving minor variations of the same purpose... sometimes within the same file!
My theory is that this is due to a kind of "tunnel vision" these models have as they execute on a given task, because engineers new to a company do the same thing until they learn the "lay of the land" and figure out that similar …
It connects the need for semantic search to the real-world problem of LLM coding agents writing redundant code due to context tunnel vision.
Generated by an LLM from the article and thread. Check the sources before relying on it.