Codebase RAG: Building Context Graphs for AI Coding Tools

Key Takeaways

• Codebase RAG built on chunk-and-embed alone underperforms because source code is a graph of definitions, calls, imports, and inheritance — splitting it into flat text windows destroys exactly the structure a coding agent needs
• The fix is a context graph: parse every file into an AST with Tree-sitter, resolve symbols into nodes and relationships, store the structure in a graph database such as Neo4j and the semantics in a vector index such as pgvector
• Retrieval becomes two steps, not one — semantic search finds the entry point, graph traversal expands it to callers, callees, type definitions, and tests, so the model gets a closed set of relevant code instead of the top-k nearest chunks
• Use a code-specialised embedding model: voyage-code-3 outperformed OpenAI-v3-large by an average of 13.80% across 32 code retrieval datasets (Voyage AI)
• Precision beats volume. Stanford and UW's Lost in the Middle study found retrieval accuracy degrades sharply when the relevant passage sits mid-context (Liu et al., TACL 2024) — a tight, structurally complete context outperforms a large, loosely relevant one

Point a standard RAG pipeline at a repository and it will answer questions about your README beautifully and questions about your code badly. The retrieval looks healthy — cosine scores are high, chunks come back fast — but the agent keeps proposing changes that break callers it never saw, because nothing in a flat vector index knows that PaymentService.refund() has nine call sites across four modules.

Codebase RAG is a different retrieval problem from document RAG, and it needs a different index.

Codebase RAG implementation works by combining a structural graph of the repository with semantic vector search over its contents. You parse source files into ASTs, extract symbols and their relationships into a queryable context graph, embed the code with a code-specialised model, then retrieve by finding a semantic entry point and traversing the graph outward to everything structurally connected to it.

Why Chunk-and-Embed Fails on Source Code

Standard RAG chunking strategies assume prose: text that degrades gracefully when split, where a 512-token window usually contains a complete thought. Source code violates that assumption at every level. A function body is meaningless without its type definitions, a class is incomplete without its base class, and a change is unsafe without its callers.

Three specific failure modes show up in every naive implementation we have reviewed:

Split-symbol retrieval. A fixed-size splitter cuts a 200-line class in half. The retriever returns the half containing the matching method name; the half containing the constructor, the field types, and the imports never arrives. The model hallucinates the missing half.

Lexical near-misses. Codebases are full of near-identical text. getUser, getUserById, getUserByEmail, and their four test doubles all embed to nearly the same vector. Top-k similarity cannot separate them, so the agent picks whichever scored 0.003 higher.

Invisible dependencies. The most expensive failure. Nothing about refund()'s text tells the retriever that a billing cron job calls it. That edge exists in the call graph, not in the token stream, so an embedding index cannot represent it at any chunk size.

The pattern here is consistent: the information the agent needs is real and present in the repository, but it is encoded in relationships between files rather than in the content of any single one. Retrieval that only reads content will never find it.

What a Codebase Context Graph Actually Contains

A codebase context graph is a property graph whose nodes are code entities — repository, file, module, class, function, method, variable, test — and whose edges are the relationships a compiler already understands: DEFINES, CALLS, IMPORTS, INHERITS_FROM, IMPLEMENTS, RETURNS_TYPE, TESTED_BY. Every node carries its source span so you can fetch exact text on demand.

The open-source Potpie project is the clearest public reference for this shape: it converts a repository into a Neo4j-backed knowledge graph and runs agents on top of semantic search, call-graph traversal, and impact analysis. Prodinit built the same core structure into DevOS, our internal codebase-understanding system, and the node/edge schema converged on almost exactly that list.

Two design decisions matter more than the rest.

Keep the graph and the vectors separate but joined. The graph database answers structural questions (what calls this, what does this inherit from, which tests cover it). The vector index answers semantic ones (where is the retry logic). Each node's ID is the join key. Trying to force structure into vector metadata filters — or semantics into graph properties — produces a system that is bad at both.

Store spans, not text. Nodes hold file_path, start_line, end_line, and a content hash. Source text is read from disk at assembly time. This keeps the graph small, makes incremental updates cheap, and means the context you hand the model is always the code as it exists right now rather than a stale copy from the last index run. It is the same principle as the file system as context pattern: the repository stays the source of truth, and the index is a navigational layer over it.

Building the Graph: Parse, Resolve, Embed

The ingestion pipeline has three stages and should be idempotent per file, because you will run it thousands of times. Parse each file into an AST, resolve symbol references into edges, then embed each node's code with a code-specialised model. Budget for the resolve stage — it is where the real engineering sits.

Stage 1 — Parse with Tree-sitter. Tree-sitter gives you concrete syntax trees for dozens of languages behind one API, parses files that do not compile (essential — real repositories are often mid-refactor), and supports incremental reparsing. Walk the tree and emit a node per declaration with its span, signature, docstring, and decorators.

Stage 2 — Resolve references into edges. Parsing gives you nodes; resolution gives you the graph. For each call expression, import, and type annotation, resolve the identifier to the node it refers to using scope rules and the import map. Get this right and your graph is useful; get it wrong and you have an expensive list of files.

Two rules keep resolution honest in production:

• Resolve within a language, link across languages by convention. A Python service calling a TypeScript endpoint will not resolve through the AST. Match those edges on route strings or schema names and label them INFERRED so downstream consumers can weight them lower.
• Never drop an unresolved reference. Write it as a dangling edge with the raw identifier. Unresolved references are the single best signal for measuring graph quality over time, and silently discarding them hides regressions.

Stage 3 — Embed with a code model, not a text model. This is the cheapest accuracy win available. voyage-code-3 outperformed OpenAI-v3-large by an average of 13.80% and CodeSage-large by 16.81% across a suite of 32 code retrieval datasets (Voyage AI, December 2024). Embed at node granularity — one vector per function or method rather than per arbitrary window — so the retrieval unit and the graph unit are the same thing.

Enrich the embedded text before you send it. A function embedded as bare source competes poorly against a function embedded as file path + class name + signature + docstring + body, because the surrounding identifiers are often where the searchable intent actually lives.

Codebase RAG Retrieval: Semantic Entry, Structural Expansion

Codebase RAG retrieval runs in two stages: a vector search that finds the most semantically relevant entry-point nodes, then a bounded graph traversal that expands those nodes into a structurally complete context. The traversal is what separates this from ordinary RAG — it is the step that finds the nine callers a similarity search would never surface.

A query like "why does refund fail for partial captures" resolves like this:

1. Entry point. Vector search over node embeddings returns the top 5–10 candidates — likely PaymentService.refund, RefundValidator.validate, and a test.
2. Expansion. From each candidate, traverse outward with a hop budget: callees at depth 1–2 (what this code depends on), callers at depth 1 (what breaks if it changes), type definitions for every parameter and return type, and linked tests.
3. Assembly. Deduplicate by node ID, order by dependency so definitions precede usage, and fetch source text by span.
4. Budget. Trim to the context budget by dropping the lowest-scoring leaf nodes first, never the structural spine — a callee whose definition is missing is worse than no callee at all.

Bound the traversal hard. When retrieval quality is poor, the instinct is to raise the hop budget — but depth 3 on a well-connected repository pulls in most of it, and a context that large makes quality worse, not better. Lost in the Middle found accuracy on retrieval tasks degrades significantly when the relevant passage sits in the middle of a long context rather than at either end (Liu et al., TACL 2024). A 6k-token context containing exactly the right call chain beats a 100k-token context containing it somewhere.

Rank the expanded set before assembly. Combine semantic score, graph distance from the entry point, and a recency signal from git history — code changed in the last 30 days is disproportionately likely to be relevant to a question someone is asking today.

Keeping the Graph Fresh

A context graph that is a day stale is worse than no graph, because the agent trusts it. Incremental updates on commit are mandatory: diff the changed files, reparse only those, delete and rewrite their nodes, then repair only the edges that touched them. Full re-indexing is for schema migrations, not for Tuesdays.

The practical implementation is a webhook on push that computes changed paths, reparses each one with Tree-sitter, and replaces that file's subgraph transactionally. Two details make the difference between a system that stays correct and one that drifts:

• Re-resolve inbound edges, not just outbound ones. When a function is deleted or renamed, its callers still point at it. Those inbound edges must be re-resolved or explicitly marked dangling, or the graph quietly accumulates references to code that no longer exists.
• Key embeddings by content hash. Most commits touch files without changing most of the functions inside them. Hashing each node's source and skipping unchanged hashes means a typical commit re-embeds a handful of nodes rather than every node in every file it touched — which is where most of the embedding spend on an active repository would otherwise go.

Track two metrics continuously: the unresolved-reference rate (graph quality) and retrieval recall against a golden set of 30–50 annotated repository questions with known correct files (retrieval quality). Both should be part of CI, the same way you would gate any other production RAG pipeline. Without them, quality regressions in the index are invisible until an agent ships a bad patch.

The Stack, End to End

A production codebase RAG implementation is five components, none exotic: Tree-sitter for parsing, a graph store for structure, a vector index for semantics, a code embedding model, and an assembly layer that enforces the context budget. The choices within each are mostly interchangeable — the architecture is what matters.

The smallest version that works is worth building first: Tree-sitter parsing, nodes and edges in PostgreSQL, pgvector for embeddings, one-hop expansion. Prodinit shipped exactly that shape as the first working version of DevOS, then added cross-language edge inference and git-recency ranking once the golden set showed where one-hop expansion was still missing context. Add Neo4j when traversal depth, not correctness, becomes the bottleneck.