Retrieval-augmented generation for code search in large monorepos

By Femi Adeyemo · 16 July 20260 views

Searching a 2-million-line monorepo is a different problem from searching a documentation site. Keywords are ambiguous: render appears in 400 files spanning a React frontend, a PDF generation library, and a WebGL renderer. Symbol names are often short and collision-prone. The most important information is often not in the file you are looking at but in the call sites, type definitions, and tests that surround it.

Existing solutions are inadequate at scale. grep is fast but has no semantic understanding. Language Server Protocol symbol search is precise but does not understand intent. Embedding-based search systems built for prose documents do not handle code well because code's semantics live in structure and types, not just token co-occurrence.

RAG for code search addresses this by combining embedding-based retrieval with an LLM that can reason across multiple retrieved chunks. The result is a system that can answer questions like "where is the discount calculation logic and what edge cases does it handle?" in seconds, with citations to specific files and line numbers.

This article builds the full pipeline: code chunking, embedding generation, a pgvector index, and a retrieval-augmented answer layer using Claude.

Code chunking strategy

Generic text chunkers — splitting by character count or sentence boundaries — produce terrible code chunks. They cut in the middle of functions, separate a method from its class context, and discard the import statements that define what symbols mean.

A language-aware chunker splits at logical boundaries and preserves context:

import ast
from dataclasses import dataclass
from pathlib import Path
from typing import Iterator

@dataclass
class CodeChunk:
    file_path: str
    start_line: int
    end_line: int
    content: str
    chunk_type: str  # "function", "class", "method", "module_block"
    symbol_name: str | None
    imports: str  # Import section at the top of the file


def chunk_python_file(file_path: Path) -> Iterator[CodeChunk]:
    source = file_path.read_text(encoding="utf-8", errors="replace")
    
    try:
        tree = ast.parse(source)
    except SyntaxError:
        # Fall back to line-based chunking for unparseable files
        yield from chunk_by_lines(file_path, source)
        return
    
    lines = source.splitlines()
    imports = extract_imports(tree, lines)
    
    for node in ast.walk(tree):
        if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
            if not _is_nested(node, tree):
                yield _make_chunk(node, lines, imports, file_path, "function")
        
        elif isinstance(node, ast.ClassDef):
            # Yield the class definition (without method bodies) as one chunk
            yield _make_class_header_chunk(node, lines, imports, file_path)
            
            # Yield each method as a separate chunk with class context
            for item in node.body:
                if isinstance(item, (ast.FunctionDef, ast.AsyncFunctionDef)):
                    class_context = f"# In class: {node.name}\n"
                    yield _make_chunk(
                        item, lines, class_context + imports, file_path, "method",
                        parent_name=node.name
                    )


def extract_imports(tree: ast.Module, lines: list[str]) -> str:
    import_lines = []
    for node in tree.body:
        if isinstance(node, (ast.Import, ast.ImportFrom)):
            start = node.lineno - 1
            end = node.end_lineno
            import_lines.extend(lines[start:end])
    return "\n".join(import_lines)


def _make_chunk(
    node: ast.AST,
    lines: list[str],
    context_prefix: str,
    file_path: Path,
    chunk_type: str,
    parent_name: str | None = None
) -> CodeChunk:
    start = node.lineno - 1
    end = node.end_lineno
    content = "\n".join(lines[start:end])
    symbol = getattr(node, "name", None)
    
    # Prepend import context so embedding captures dependency information
    full_content = f"# File: {file_path}\n{context_prefix}\n\n{content}"
    
    return CodeChunk(
        file_path=str(file_path),
        start_line=node.lineno,
        end_line=node.end_lineno,
        content=full_content,
        chunk_type=chunk_type,
        symbol_name=f"{parent_name}.{symbol}" if parent_name else symbol,
        imports=context_prefix,
    )

For TypeScript, use the TypeScript compiler API to walk the AST in the same way. The key invariant is that every chunk is self-contained: it includes enough import context that an embedding model (or a language model reading it later) can understand what the symbols mean without seeing the rest of the file.

Chunks that are too short (under 5 lines) carry no signal and pollute retrieval results. Chunks over 150 lines often contain too much to embed meaningfully. Filter both extremes:

MIN_CHUNK_LINES = 5
MAX_CHUNK_LINES = 150

def filter_chunks(chunks: list[CodeChunk]) -> list[CodeChunk]:
    return [
        c for c in chunks
        if MIN_CHUNK_LINES <= (c.end_line - c.start_line) <= MAX_CHUNK_LINES
    ]

Generating code-aware embeddings

Code embeddings work better with models trained on code than with general-purpose text embeddings. text-embedding-3-large from OpenAI performs well on code; voyage-code-2 from Voyage AI is specifically optimised for it. Both produce 1536-dimensional vectors.

import openai
from concurrent.futures import ThreadPoolExecutor, as_completed

embedding_client = openai.OpenAI()

def embed_chunks(chunks: list[CodeChunk], batch_size: int = 100) -> list[list[float]]:
    """Embed code chunks in batches and return vectors."""
    all_embeddings = []
    
    for i in range(0, len(chunks), batch_size):
        batch = chunks[i:i + batch_size]
        texts = [c.content for c in batch]
        
        response = embedding_client.embeddings.create(
            model="text-embedding-3-large",
            input=texts,
            dimensions=1536
        )
        
        batch_embeddings = [item.embedding for item in response.data]
        all_embeddings.extend(batch_embeddings)
        
        print(f"Embedded {min(i + batch_size, len(chunks))}/{len(chunks)} chunks")
    
    return all_embeddings

For a monorepo with 50,000 chunks, embedding generation takes about 10 minutes and costs around $1.50 at current OpenAI pricing. Run it once, then update incrementally as files change.

Storing chunks in pgvector

PostgreSQL with the pgvector extension is the right choice for production code search: it runs on your existing infrastructure, supports filtered queries (filter by file path, language, or symbol type before the vector search), and handles the dataset size of a typical monorepo without specialist vector database expertise.

Set up the schema:

CREATE EXTENSION IF NOT EXISTS vector;

CREATE TABLE code_chunks (
    id BIGSERIAL PRIMARY KEY,
    repo_id TEXT NOT NULL,
    file_path TEXT NOT NULL,
    start_line INT NOT NULL,
    end_line INT NOT NULL,
    chunk_type TEXT NOT NULL,  -- function, method, class, module_block
    symbol_name TEXT,
    language TEXT NOT NULL,
    content TEXT NOT NULL,
    imports TEXT,
    embedding vector(1536),
    indexed_at TIMESTAMPTZ DEFAULT NOW(),
    
    UNIQUE(repo_id, file_path, start_line, end_line)
);

-- IVFFlat index for approximate nearest neighbour search
-- Use 100 lists for datasets up to 1M chunks
CREATE INDEX ON code_chunks 
    USING ivfflat (embedding vector_cosine_ops) 
    WITH (lists = 100);

-- Covering index for common filter patterns
CREATE INDEX ON code_chunks (repo_id, language, chunk_type);

Insert chunks with their embeddings:

import psycopg2
from psycopg2.extras import execute_values

def store_chunks(
    chunks: list[CodeChunk],
    embeddings: list[list[float]],
    repo_id: str,
    language: str,
    conn_string: str,
):
    conn = psycopg2.connect(conn_string)
    
    rows = [
        (
            repo_id,
            chunk.file_path,
            chunk.start_line,
            chunk.end_line,
            chunk.chunk_type,
            chunk.symbol_name,
            language,
            chunk.content,
            chunk.imports,
            embedding,
        )
        for chunk, embedding in zip(chunks, embeddings)
    ]
    
    with conn.cursor() as cur:
        execute_values(
            cur,
            """
            INSERT INTO code_chunks 
                (repo_id, file_path, start_line, end_line, chunk_type, 
                 symbol_name, language, content, imports, embedding)
            VALUES %s
            ON CONFLICT (repo_id, file_path, start_line, end_line) 
            DO UPDATE SET
                content = EXCLUDED.content,
                embedding = EXCLUDED.embedding,
                indexed_at = NOW()
            """,
            rows,
            template="(%s, %s, %s, %s, %s, %s, %s, %s, %s, %s::vector)"
        )
    
    conn.commit()
    conn.close()

Retrieval with pre-filtering

A raw vector search across all chunks returns the most semantically similar code, but often you want to scope the search. If a developer asks about "the auth module," you do not want results from the logging module. Pre-filter by metadata before the vector search:

def search_code(
    query: str,
    repo_id: str,
    conn_string: str,
    language: str | None = None,
    chunk_types: list[str] | None = None,
    file_path_prefix: str | None = None,
    top_k: int = 10,
) -> list[dict]:
    # Embed the query
    query_embedding = embedding_client.embeddings.create(
        model="text-embedding-3-large",
        input=[query],
        dimensions=1536
    ).data[0].embedding
    
    # Build filtered query
    filters = ["repo_id = %s"]
    params = [repo_id]
    
    if language:
        filters.append("language = %s")
        params.append(language)
    
    if chunk_types:
        filters.append("chunk_type = ANY(%s)")
        params.append(chunk_types)
    
    if file_path_prefix:
        filters.append("file_path LIKE %s")
        params.append(f"{file_path_prefix}%")
    
    where_clause = " AND ".join(filters)
    params.extend([query_embedding, top_k])
    
    sql = f"""
    SELECT 
        file_path,
        start_line,
        end_line,
        chunk_type,
        symbol_name,
        content,
        1 - (embedding <=> %s::vector) AS similarity
    FROM code_chunks
    WHERE {where_clause}
    ORDER BY embedding <=> %s::vector
    LIMIT %s
    """
    
    # Note: params need reordering to match SQL placeholders
    conn = psycopg2.connect(conn_string)
    with conn.cursor() as cur:
        cur.execute(sql, params[:-2] + [query_embedding, query_embedding, top_k])
        rows = cur.fetchall()
    conn.close()
    
    return [
        {
            "file_path": row[0],
            "start_line": row[1],
            "end_line": row[2],
            "chunk_type": row[3],
            "symbol_name": row[4],
            "content": row[5],
            "similarity": row[6],
        }
        for row in rows
    ]

For production use, set ivfflat.probes to 10 before vector queries for better recall at a modest performance cost:

SET ivfflat.probes = 10;

The answer layer with Claude

Retrieved chunks are raw code. The developer's question needs a synthesised answer that explains what the code does, how the pieces fit together, and where to look for more context. Claude handles this well when you structure the prompt to include all retrieved context:

import anthropic

claude_client = anthropic.Anthropic()

def answer_code_question(
    question: str,
    retrieved_chunks: list[dict],
    repo_name: str,
) -> dict:
    context = _format_retrieved_chunks(retrieved_chunks)
    
    response = claude_client.messages.create(
        model="claude-opus-4-5",
        max_tokens=2048,
        system=f"""You are an expert software engineer with deep knowledge of the {repo_name} codebase.
When answering questions about code, you:
- Cite specific files and line numbers using the format [file:line]
- Explain not just what the code does but why it is structured that way
- Note any edge cases or gotchas visible in the retrieved code
- Suggest where to look for more context when the retrieved chunks are incomplete
- Never make up code that was not in the retrieved context""",
        messages=[{
            "role": "user",
            "content": f"""Question: {question}

Relevant code from the codebase:

{context}

Answer the question based only on the code above. Cite file paths and line numbers."""
        }]
    )
    
    return {
        "question": question,
        "answer": response.content[0].text,
        "sources": [
            {
                "file": c["file_path"],
                "lines": f"{c['start_line']}-{c['end_line']}",
                "symbol": c.get("symbol_name"),
                "similarity": round(c["similarity"], 3),
            }
            for c in retrieved_chunks
        ]
    }

def _format_retrieved_chunks(chunks: list[dict]) -> str:
    sections = []
    for chunk in chunks:
        header = f"### {chunk['file_path']} (lines {chunk['start_line']}-{chunk['end_line']})"
        if chunk.get("symbol_name"):
            header += f" — {chunk['symbol_name']}"
        sections.append(f"{header}\n\n```\n{chunk['content']}\n```")
    return "\n\n".join(sections)

Keeping the index current

A code search index that goes stale is worse than no index: developers trust stale results and make wrong decisions. Incremental indexing based on git changes is essential.

import subprocess

def get_changed_files_since(commit: str) -> list[str]:
    result = subprocess.run(
        ["git", "diff", "--name-only", commit, "HEAD"],
        capture_output=True, text=True, check=True
    )
    return [f for f in result.stdout.strip().split("\n") if f]

def incremental_update(
    repo_id: str,
    repo_root: Path,
    last_indexed_commit: str,
    conn_string: str,
):
    changed_files = get_changed_files_since(last_indexed_commit)
    
    python_files = [
        repo_root / f for f in changed_files
        if f.endswith(".py") and (repo_root / f).exists()
    ]
    
    deleted_files = [
        f for f in changed_files
        if not (repo_root / f).exists()
    ]
    
    # Remove chunks for deleted/modified files
    all_to_reindex = changed_files
    conn = psycopg2.connect(conn_string)
    with conn.cursor() as cur:
        cur.execute(
            "DELETE FROM code_chunks WHERE repo_id = %s AND file_path = ANY(%s)",
            (repo_id, all_to_reindex)
        )
    conn.commit()
    conn.close()
    
    # Re-index modified files
    for file_path in python_files:
        chunks = list(chunk_python_file(file_path))
        chunks = filter_chunks(chunks)
        if chunks:
            embeddings = embed_chunks(chunks)
            store_chunks(chunks, embeddings, repo_id, "python", conn_string)
    
    current_commit = subprocess.run(
        ["git", "rev-parse", "HEAD"],
        capture_output=True, text=True
    ).stdout.strip()
    
    return current_commit

Run this as a post-receive git hook or a CI job triggered by pushes to main. At typical commit rates, incremental updates take under a minute for repositories with thousands of files.

Putting the search API together

Expose the pipeline through a simple HTTP API:

from fastapi import FastAPI
from pydantic import BaseModel

app = FastAPI()

class SearchRequest(BaseModel):
    question: str
    repo_id: str
    language: str | None = None
    file_prefix: str | None = None

@app.post("/search")
async def search(req: SearchRequest):
    chunks = search_code(
        query=req.question,
        repo_id=req.repo_id,
        conn_string=CONN_STRING,
        language=req.language,
        file_path_prefix=req.file_prefix,
        top_k=8,
    )
    
    # Filter to similarity threshold — below 0.65 is usually noise
    chunks = [c for c in chunks if c["similarity"] >= 0.65]
    
    if not chunks:
        return {"answer": "No relevant code found for this question.", "sources": []}
    
    result = answer_code_question(req.question, chunks, req.repo_id)
    return result

A full query — embedding the question, searching pgvector, and generating an answer with Claude — completes in 2-4 seconds on typical hardware. Most of that time is the Claude API call; the pgvector search itself is under 100 ms for datasets of 100,000 chunks.

The combination of AST-aware chunking, contextual import preservation, and a reasoning layer that synthesises multiple retrieved chunks produces meaningfully better results than keyword search for the kind of questions engineers actually ask when navigating an unfamiliar codebase.

Comments

No comments yet. Be the first!

Sign in to leave a comment.