CodeRAG: AST-Parsed Codebase RAG with Hybrid Search and Precise Citations

Why Standard RAG Fails on Code

Standard document RAG chunks text at fixed character offsets. For prose, this works reasonably well — the semantic content of a paragraph usually survives arbitrary splitting.

For code, it's a disaster. A 512-token chunk that starts mid-function has lost its function signature. A chunk that ends mid-if-statement has lost its else branch. The LLM receives context that's syntactically incomplete and semantically meaningless.

CodeRAG solves this with AST-aware chunking.

AST-Aware Chunking with Tree-sitter

Tree-sitter is a parser generator that produces concrete syntax trees for dozens of programming languages. Instead of splitting on character count, we split on syntactic boundaries:

import tree_sitter_python as tspython
from tree_sitter import Language, Parser

PY_LANGUAGE = Language(tspython.language())
parser = Parser(PY_LANGUAGE)

def extract_python_chunks(source: str, filepath: str) -> list[CodeChunk]:
    """Extract semantically complete code units from a Python file."""
    tree = parser.parse(source.encode())
    chunks = []
    
    def visit(node, parent_name=""):
        if node.type in ("function_definition", "class_definition"):
            chunk_text = source[node.start_byte:node.end_byte]
            chunks.append(CodeChunk(
                content=chunk_text,
                filepath=filepath,
                start_line=node.start_point[0] + 1,
                end_line=node.end_point[0] + 1,
                node_type=node.type,
                name=get_node_name(node, source),
            ))
        for child in node.children:
            visit(child, parent_name)
    
    visit(tree.root_node)
    return chunks

Every chunk is a complete syntactic unit: a full function, a complete class, a module-level statement. The LLM always receives syntactically valid, semantically complete code.

Hybrid Retrieval: Dense + Sparse

Dense embedding search (cosine similarity on vectors) is excellent at semantic retrieval — "find functions that do authentication" returns auth-related functions even if they don't contain the word "authentication". But it misses exact identifier matches: searching for UserSessionManager in embedding space might not return the class if the embedding space squashes exact names.

Sparse BM25 search (keyword-based inverted index) is perfect for exact identifier lookup but blind to semantic meaning.

Hybrid retrieval combines both:

def hybrid_retrieve(query: str, k: int = 8) -> list[CodeChunk]:
    # Dense retrieval
    query_embedding = embed(query)
    dense_results = pgvector_search(query_embedding, k=k*2)
    
    # Sparse retrieval
    sparse_results = bm25_search(query, k=k*2)
    
    # Reciprocal Rank Fusion
    scores = {}
    for rank, chunk in enumerate(dense_results):
        scores[chunk.id] = scores.get(chunk.id, 0) + 1 / (60 + rank)
    for rank, chunk in enumerate(sparse_results):
        scores[chunk.id] = scores.get(chunk.id, 0) + 1 / (60 + rank)
    
    # Return top-k by fused score
    ranked = sorted(scores.items(), key=lambda x: x[1], reverse=True)
    return [get_chunk(id) for id, _ in ranked[:k]]

Reciprocal Rank Fusion (RRF) is the score fusion method — it's rank-based, so it doesn't require normalizing incompatible score scales.

Precise File and Line Citations

Every retrieved chunk carries its file path, start line, and end line from the AST parsing step. The LLM generation prompt includes this metadata:

def build_context(chunks: list[CodeChunk]) -> str:
    context_parts = []
    for chunk in chunks:
        context_parts.append(
            f"[{chunk.filepath}:{chunk.start_line}-{chunk.end_line}]\n"
            f"```{chunk.language}\n{chunk.content}\n```"
        )
    return "\n\n".join(context_parts)

prompt = f"""
You are a code analysis assistant. Answer questions about the codebase using ONLY the code snippets below.
For every claim, cite the file and line range in the format [filepath:start-end].

Code Context:
{build_context(retrieved_chunks)}

Question: {query}
"""

The result: every answer includes precise, verifiable source references. Not "I think this is in the auth module" — but "this is in src/auth/session.py:142-167."

pgvector Setup

CREATE EXTENSION IF NOT EXISTS vector;

CREATE TABLE code_chunks (
    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
    filepath TEXT NOT NULL,
    start_line INTEGER NOT NULL,
    end_line INTEGER NOT NULL,
    node_type TEXT NOT NULL,
    content TEXT NOT NULL,
    embedding vector(1536)  -- OpenAI text-embedding-3-small dimension
);

CREATE INDEX ON code_chunks USING ivfflat (embedding vector_cosine_ops)
    WITH (lists = 100);

Evaluation Results

Across a 50-question test set on a medium-sized TypeScript codebase:

Metric	Character Chunking	AST Chunking
Retrieval Precision @5	0.61	0.84
Citation Accuracy	N/A	0.91
Syntactically Valid Chunks	71%	100%

AST chunking improved retrieval precision by 38%. Every chunk is syntactically valid. Citation accuracy at 91% means 91% of generated citations are verifiable in the source file.

GitHub: github.com/karthikrshet/coderag