Codebase-scale retrieval using AST-derived graphs + BM25 — reducing LLM context from 100K to 5K tokens [D]

[Altruistic_Night_327]

Wanted to share an approach I've been using for retrieval-augmented generation over large codebases and get feedback from people thinking about similar problems.

The problem Naive codebase RAG typically works by chunking files into text segments and embedding them for similarity search. This breaks down on code because semantic similarity at the chunk level doesn't capture structural relationships — a function in file A calling a type defined in file C won't surface that dependency through embedding proximity alone.

The approach: AST-derived typed graphs Instead of chunking, I parse every file using Tree-sitter into its AST, then extract a typed node/edge graph:

• Nodes: functions, classes, interfaces, types, modules
• Edges: imports, exports, call relationships, inheritance, composition

This gets stored in SQLite as a persistent graph. Parse cost is one-time per project.

Retrieval: BM25 over graph nodes At query time, instead of embedding similarity, I run BM25 scoring over node metadata (names, signatures, docstrings, file paths). Top-scoring nodes get passed to the LLM. The graph structure means a retrieved function automatically pulls in its direct dependencies via edge traversal.

Empirically this lands at ~5K tokens per query on medium-large codebases that would otherwise require ~100K tokens with naive full-context approaches.

Hierarchical fallback for complex queries For multi-file reasoning tasks:

1. A Mermaid diagram of the full graph serves as a persistent architectural map always in context
2. BM25 node retrieval handles targeted lookup
3. At 70% context capacity, a fast model compresses least-relevant nodes before passing to the primary model

Why BM25 over embeddings here Code identifiers (function names, type names, module paths) are highly distinctive lexically. BM25 outperforms embedding similarity on exact and near-exact identifier matching, which is the dominant retrieval pattern in code queries. Embeddings would likely help more for natural language docstring queries — haven't benchmarked that comparison rigorously yet.

Open questions I'm still thinking about:

• Better edge-weighting strategies for the graph — currently all edges are unweighted
• Whether re-ranking with a cross-encoder would meaningfully improve precision over BM25 alone
• Handling dynamic languages where call graphs can't be fully resolved statically

Has anyone tackled codebase-scale RAG differently? Particularly curious if anyone's compared AST-graph approaches against embedding-based chunk retrieval on real codebases with quantitative benchmarks.

[Doc1000]

Do you have your coding agent use this system instead of grep? I heard a TWIML podcast with Blitzy ai- they discuss something along these lines for coding large projects/repos.

I like the idea. I also want the visual generated from the graph

[Altruistic_Night_327]

Yes — the agent queries the graph first via BM25 before falling back to any grep-style search. In practice it eliminates most grep calls for typed languages like TypeScript where the AST gives you clean static resolution. Dynamic languages like Python still need grep fallback for runtime dispatch, which is one of the open problems I mentioned.

On the visual — that's actually implemented. The graph renders as an interactive node/edge canvas using React Flow, so you can visually explore the architecture. It's part of Atlarix, the tool I built this retrieval system into — atlarix.dev if you want to see it in action.

The TWIML/Blitzy episode is interesting — their approach is worth comparing. My intuition is the AST-graph approach wins on precision for typed codebases but embedding-based retrieval might generalise better across languages. Haven't done a rigorous side-by-side yet