Retrieval-augmented generation for code search in large monorepos
Why code search is harder than document search
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.