11.3 Retrieval-Augmented Code Generation

1. Dense Retrieval with Code Embeddings

When a developer types "find the retry logic for API calls" and the system surfaces a function named _backoff_request from a file they have never opened, keyword matching cannot explain the result. Dense retrieval makes this leap by encoding both the query and every candidate chunk as vectors in a shared embedding space, then retrieving the \(k\) chunks whose vectors lie closest to the query vector. The quality of this approach depends entirely on the embedding model: it must map semantically similar code and natural-language descriptions to nearby points.

Dense retrieval converts natural-language queries and code fragments into fixed-length numerical vectors (embeddings), reducing semantic similarity to geometric proximity. A developer describes what they need in plain English ("find the retry logic for API calls") and retrieves relevant code even when no keywords overlap. An encoder neural network maps text to a point in high-dimensional space. A nearest-neighbor search then returns the \(k\) closest stored points. Use dense retrieval when queries describe intent or behavior. Switch to keyword-based methods (BM25, grep) when queries contain exact identifiers, error strings, or class names that must match literally.

Choosing an Encoder

Once you have committed to dense retrieval for an intent-based query, the critical engineering decision becomes which encoder to use, because the embedding model determines the geometry of that shared space.

General-purpose text embedding models (e.g., all-MiniLM-L6-v2) work well on code because natural-language tokens in docstrings and variable names carry semantic signal, but code-specialized models perform better. Models trained on the CodeSearchNet benchmark (a dataset of six programming languages pairing functions with their docstrings, used to evaluate code search quality), such as microsoft/unixcoder-base and Salesforce/codet5p-110m-embedding, learn to align code with its natural-language description during pre-training (as of 2025, newer code embedding models such as Voyage AI's voyage-code-3 and OpenAI's text-embedding-3-large significantly outperform these earlier models on retrieval benchmarks, though the architectural principle of code-specialized pretraining remains the same). The cosine similarity between a query "compute attention scores" and a function implementing scaled dot-product attention should be high, even though the query and the code share few surface tokens.

The similarity between a query embedding \(\mathbf{q}\) and a chunk embedding \(\mathbf{c}\) is the cosine similarity:

For normalized embeddings (which most sentence-transformer models produce), this simplifies to the dot product \(\mathbf{q} \cdot \mathbf{c}\), making retrieval equivalent to a maximum inner product search (MIPS) problem (finding the stored vector whose dot product with the query vector is largest). In short: retrieval-augmented code generation turns "find the right file" from a human navigation problem into a geometry problem, where the best context is simply the nearest point.

import numpy as np
from sentence_transformers import SentenceTransformer

class DenseRetriever:
    """Embedding-based code retrieval using sentence-transformers."""

    def __init__(self, model_name: str = "all-MiniLM-L6-v2"):
        self.model = SentenceTransformer(model_name)
        self.chunks: list[dict] = []         # chunk metadata
        self.embeddings: np.ndarray | None = None  # (n_chunks, embed_dim)

    def index(self, chunks: list[dict]):
        """Embed and index a list of chunks.

        Each chunk dict must have 'id', 'content', and optionally
        'llm_summary' (preferred for embedding if available).
        """
        self.chunks = chunks

        # Prefer LLM summaries for embedding (better NL-to-NL matching)
        texts = []
        for c in chunks:
            if c.get("llm_summary"):
                texts.append(c["llm_summary"])
            elif c.get("signature"):
                texts.append(c["signature"])
            else:
                texts.append(c["content"][:2000])  # truncate long chunks

        self.embeddings = self.model.encode(
            texts, normalize_embeddings=True, show_progress_bar=True,
            batch_size=64,
        )

    def search(self, query: str, top_k: int = 20) -> list[tuple[dict, float]]:
        """Return top-k chunks ranked by cosine similarity to the query."""
        query_embedding = self.model.encode(
            [query], normalize_embeddings=True
        )[0]

        # Dot product = cosine similarity for normalized vectors
        scores = self.embeddings @ query_embedding
        top_indices = np.argsort(scores)[::-1][:top_k]

        return [(self.chunks[i], float(scores[i])) for i in top_indices]

2. Sparse Retrieval with BM25

Dense retrieval excels at semantic matching ("find functions related to authentication") but struggles with exact identifier matching ("find all uses of validate_jwt_token"). A query containing a specific function name, class name, or error message needs lexical matching, not semantic similarity. This is where BM25 shines.

BM25 is a probabilistic ranking function that extends term frequency-inverse document frequency (TF-IDF) with term frequency saturation (diminishing returns as a term repeats more often in a document) and document length normalization. For a query \(Q\) with terms \(q_1, \ldots, q_m\) and a document \(D\), the BM25 score is:

where \(f(q_i, D)\) is the term frequency of \(q_i\) in \(D\), \(|D|\) is the document length, \(\text{avgdl}\) is the average document length across the corpus, \(k_1\) controls term frequency saturation (typically 1.2-2.0), and \(b\) controls length normalization (typically 0.75). The IDF (inverse document frequency) component is:

where \(N\) is the total number of documents and \(n(q_i)\) is the number of documents containing term \(q_i\).

from rank_bm25 import BM25Okapi
import re

class SparseRetriever:
    """BM25-based keyword retrieval for code chunks."""

    def __init__(self):
        self.chunks: list[dict] = []
        self.bm25: BM25Okapi | None = None

    @staticmethod
    def tokenize_code(text: str) -> list[str]:
        """Tokenize code for BM25: split on camelCase, snake_case, punctuation."""
        # Split camelCase: "validateJwtToken" -> ["validate", "Jwt", "Token"]
        text = re.sub(r'([a-z])([A-Z])', r'\1 \2', text)
        # Split on underscores and non-alphanumeric
        tokens = re.findall(r'[a-zA-Z]\w*', text.lower())
        return tokens

    def index(self, chunks: list[dict]):
        """Build the BM25 index from chunk contents."""
        self.chunks = chunks
        tokenized = [self.tokenize_code(c["content"]) for c in chunks]
        self.bm25 = BM25Okapi(tokenized)

    def search(self, query: str, top_k: int = 20) -> list[tuple[dict, float]]:
        """Return top-k chunks ranked by BM25 score."""
        query_tokens = self.tokenize_code(query)
        scores = self.bm25.get_scores(query_tokens)
        top_indices = np.argsort(scores)[::-1][:top_k]

        return [(self.chunks[i], float(scores[i])) for i in top_indices]

3. Hybrid Search and Reciprocal Rank Fusion

When a coding agent retrieves context from the wrong combination of sources, the result is not merely suboptimal: the model generates code that confidently calls functions with incorrect signatures, references modules that do not exist in the repository, or misuses internal APIs whose contracts differ subtly from their public documentation. Getting retrieval fusion right is the difference between an agent that writes code against the real codebase and one that writes against a hallucinated approximation of it.

Neither dense nor sparse retrieval dominates in all cases. Hybrid search runs both retrievers in parallel and combines their rankings. The challenge is that dense scores (cosine similarity, range \([-1, 1]\)) and BM25 scores (unbounded positive reals) are not directly comparable. We need a rank-based fusion method that is agnostic to score scales. The five-stage flow in Figure 11.3 shows where RRF sits in this pipeline. Figure 11.3.1 illustrates Hybrid RACG pipeline with RRF fusion.

RRF (Cormack et al., 2009) is the standard solution. For each document \(d\), RRF computes a fused score based on the document's rank in each retriever's result list:

where \(\mathcal{R}\) is the set of retrievers, \(\text{rank}_r(d)\) is the rank of document \(d\) in retriever \(r\)'s results (1-indexed), and \(k\) is a constant (typically 60) that prevents the top-ranked result from dominating. If a document does not appear in a retriever's top-\(n\) results, its rank is set to \(n + 1\).

def reciprocal_rank_fusion(
    ranked_lists: list[list[tuple[dict, float]]],
    k: int = 60,
    top_n: int = 20,
) -> list[tuple[dict, float]]:
    """Fuse multiple ranked lists using Reciprocal Rank Fusion (RRF).

    Args:
        ranked_lists: list of ranked results from different retrievers,
                      each entry is (chunk_dict, score)
        k: RRF constant (default 60, per the original paper)
        top_n: number of results to return

    Returns:
        Fused ranking as (chunk_dict, rrf_score) pairs, descending.
    """
    rrf_scores: dict[str, float] = {}
    chunk_lookup: dict[str, dict] = {}

    for ranked_list in ranked_lists:
        for rank, (chunk, _score) in enumerate(ranked_list, start=1):
            chunk_id = chunk["id"]
            chunk_lookup[chunk_id] = chunk
            rrf_scores[chunk_id] = rrf_scores.get(chunk_id, 0.0) + 1.0 / (k + rank)

    # Sort by fused score, descending
    sorted_ids = sorted(rrf_scores.keys(), key=lambda cid: rrf_scores[cid], reverse=True)
    return [(chunk_lookup[cid], rrf_scores[cid]) for cid in sorted_ids[:top_n]]


class HybridRetriever:
    """Combines dense and sparse retrieval with RRF fusion."""

    def __init__(self, dense: DenseRetriever, sparse: SparseRetriever):
        self.dense = dense
        self.sparse = sparse

    def search(self, query: str, top_k: int = 20,
               dense_weight: int = 50, sparse_weight: int = 50) -> list[tuple[dict, float]]:
        """Hybrid search with configurable retriever candidate pool sizes."""
        dense_results = self.dense.search(query, top_k=dense_weight)
        sparse_results = self.sparse.search(query, top_k=sparse_weight)
        return reciprocal_rank_fusion([dense_results, sparse_results], top_n=top_k)

4. Structural Re-ranking with Import Graphs

Hybrid search ranks chunks by textual relevance. But code has a structure that text does not: dependencies. If the agent is editing file A, and file A imports from file B, then chunks from file B are more likely to be relevant than chunks from an unrelated file C, even if C scores higher on textual similarity. (Structure beats semantics: a file one import edge away often outranks a file with a higher embedding score.) We exploit this by re-ranking results using the import graph from Section 11.1.

def structural_rerank(
    results: list[tuple[dict, float]],
    import_graph: dict[str, set[str]],
    focal_file: str | None = None,
    boost_factor: float = 1.5,
) -> list[tuple[dict, float]]:
    """Boost chunks from files that are structurally related to the focal file.

    A file is "structurally related" if it imports the focal file,
    is imported by the focal file, or shares an import with it.

    Args:
        results: ranked list from hybrid search
        import_graph: {module: {imported_modules}} from build_import_graph
        focal_file: the file being edited (if known)
        boost_factor: multiplicative boost for structurally related chunks

    Returns:
        Re-ranked results with structural boost applied.
    """
    if focal_file is None:
        return results

    # Normalize focal file to module path
    focal_module = focal_file.replace("/", ".").replace("\\", ".").removesuffix(".py")

    # Files imported by the focal file
    direct_imports = import_graph.get(focal_module, set())

    # Files that import the focal file (reverse edges)
    reverse_imports = {
        source for source, targets in import_graph.items()
        if focal_module in targets
    }

    # Co-imported files (share an import with focal file)
    co_imported = set()
    for source, targets in import_graph.items():
        if targets & direct_imports:  # shares at least one import
            co_imported.add(source)

    related = direct_imports | reverse_imports | co_imported

    # Apply boost
    boosted = []
    for chunk, score in results:
        chunk_module = (chunk["file_path"]
                        .replace("/", ".").replace("\\", ".")
                        .removesuffix(".py"))
        if chunk_module in related:
            boosted.append((chunk, score * boost_factor))
        else:
            boosted.append((chunk, score))

    boosted.sort(key=lambda x: x[1], reverse=True)
    return boosted

With both textual relevance and structural proximity reflected in the scores, the pipeline's remaining challenge is fitting the winning chunks into a finite context window without exceeding the model's token budget.

5. Context Packing with Position Optimization

The final step takes the re-ranked chunks and packs them into the context window. Building on the greedy knapsack from Section 11.1, we add two refinements (both visible in Listing 11.18, which sorts candidates by relevance density, meaning the ratio of retrieval score to token cost, so that compact, high-value chunks are preferred). The first is dependency inclusion: if a chunk references a symbol defined in another chunk, include that chunk too. The second is position optimization: place the most relevant chunks at the beginning and end of the context, exploiting the U-shaped attention curve (the empirical finding that LLMs recall information placed at the start or end of the context window more reliably than information in the middle) from Liu et al., 2024.

from dataclasses import dataclass, field

@dataclass
class PackedContext:
    """The final assembled context window ready for LLM consumption."""
    items: list[dict] = field(default_factory=list)
    total_tokens: int = 0
    budget: int = 0

    def to_prompt(self) -> str:
        """Render the packed context as a prompt string."""
        sections = []
        for item in self.items:
            header = f"# {item['file_path']} :: {item['name']}"
            sections.append(f"{header}\n{item['content']}")
        return "\n\n".join(sections)


def pack_context(
    ranked_results: list[tuple[dict, float]],
    budget: int,
    chunk_store: "ChunkStore",
    include_parent_signatures: bool = True,
) -> PackedContext:
    """Pack ranked chunks into a context window with dependency awareness.

    Strategy:
    1. Greedily select chunks by relevance density (score / tokens).
    2. For each selected method-level chunk, include its parent class signature.
    3. Arrange items: highest-relevance at start and end (U-shape optimization).
    """
    packed = PackedContext(budget=budget)
    selected_ids: set[str] = set()

    # Sort by relevance density
    ranked = sorted(ranked_results, key=lambda x: x[1] / max(x[0]["token_estimate"], 1),
                    reverse=True)

    for chunk, score in ranked:
        if packed.total_tokens + chunk["token_estimate"] > budget:
            continue
        if chunk["id"] in selected_ids:
            continue

        packed.items.append(chunk)
        packed.total_tokens += chunk["token_estimate"]
        selected_ids.add(chunk["id"])

        # Include parent signature if this is a method
        if include_parent_signatures and chunk.get("parent_id"):
            parent = chunk_store.get_chunk(chunk["parent_id"])
            if parent and parent["id"] not in selected_ids:
                # Use the signature summary instead of the full parent
                sig_text = parent.get("signature", "")
                if sig_text:
                    sig_item = {
                        **parent,
                        "content": sig_text,
                        "token_estimate": len(sig_text.split()) * 2,
                        "name": f"{parent['name']} (signature)",
                    }
                    if packed.total_tokens + sig_item["token_estimate"] <= budget:
                        packed.items.append(sig_item)
                        packed.total_tokens += sig_item["token_estimate"]
                        selected_ids.add(parent["id"])

    # U-shape optimization: most relevant at start and end
    if len(packed.items) >= 4:
        # Sort by original rank (approximated by order of insertion)
        # Move the second-most-relevant item to the end
        mid = len(packed.items) // 2
        first_half = packed.items[:mid]
        second_half = packed.items[mid:]
        # Reverse second half so higher-ranked items are at the end
        packed.items = first_half + second_half[::-1]

    return packed

Research Frontier: Learned Context Assembly

Hand-crafted packing heuristics (greedy knapsack, U-shape positioning) work well but leave performance on the table. RepoFormer (Wu et al., 2024) trains a small transformer to predict which combination of retrieved snippets will maximize downstream code generation quality. The model learns non-obvious interactions: including a type definition alongside a function that uses that type improved generation accuracy by 18% in their evaluation, even when the type definition scores low on standalone relevance. RLCG (Shao et al., 2024) goes further, using reinforcement learning to train the retriever end-to-end against code generation quality, bypassing relevance scoring entirely. More recently, CodeRAG-Bench (Wang et al., 2025) provides a standardized evaluation suite spanning code generation, bug fixing, and API usage tasks, revealing that retrieval quality degrades sharply when the corpus shifts from curated documentation to noisy real-world repositories. Their results show that adaptive chunk selection (dynamically choosing chunk granularity per query) closes up to 30% of the gap between fixed-granularity retrieval and oracle context, pointing toward retrieval systems that negotiate chunk boundaries at query time rather than at indexing time.

6. Evaluation Metrics for Code Retrieval

Retrieval quality requires metrics that capture both whether the right chunks appear in the results and whether the assembled context supports correct code generation. Three metrics cover the key dimensions:

Recall@k measures what fraction of the relevant chunks appear in the top-\(k\) results. If the ground truth requires chunks \(\{A, B, C\}\) and the retriever returns \(\{A, C, D, E, F\}\) at \(k=5\), then Recall@5 \(= 2/3\):

Mean Reciprocal Rank (MRR) measures how early the first relevant result appears. If the first relevant chunk is at rank 3, MRR contributes \(1/3\):

Context Utilization measures what fraction of the context budget is spent on relevant content versus noise:

def evaluate_retrieval(
    retrieved: list[dict],
    relevant_ids: set[str],
    budget: int,
) -> dict[str, float]:
    """Compute retrieval quality metrics.

    Args:
        retrieved: list of chunk dicts in ranked order
        relevant_ids: set of ground-truth relevant chunk IDs
        budget: total token budget used

    Returns:
        Dictionary of metric name -> value.
    """
    retrieved_ids = [c["id"] for c in retrieved]

    # Recall@k for various k
    metrics = {}
    for k in [5, 10, 20]:
        top_k_ids = set(retrieved_ids[:k])
        recall = len(top_k_ids & relevant_ids) / max(len(relevant_ids), 1)
        metrics[f"recall@{k}"] = recall

    # MRR: rank of first relevant result
    mrr = 0.0
    for rank, rid in enumerate(retrieved_ids, start=1):
        if rid in relevant_ids:
            mrr = 1.0 / rank
            break
    metrics["mrr"] = mrr

    # Context utilization
    relevant_tokens = sum(
        c["token_estimate"] for c in retrieved
        if c["id"] in relevant_ids
    )
    metrics["utilization"] = relevant_tokens / max(budget, 1)

    return metrics

from langchain_community.retrievers import BM25Retriever
from langchain_community.vectorstores import FAISS
from langchain.retrievers import EnsembleRetriever

bm25 = BM25Retriever.from_documents(docs, k=20)
faiss_store = FAISS.from_documents(docs, embedding_model)
faiss_retriever = faiss_store.as_retriever(search_kwargs={"k": 20})

hybrid = EnsembleRetriever(
    retrievers=[bm25, faiss_retriever],
    weights=[0.4, 0.6],  # BM25 weight, dense weight
)
results = hybrid.invoke("fix the pagination bug")

Exercises

1. (Conceptual) The RRF constant \(k = 60\) was chosen empirically by Cormack et al. Explain intuitively what happens when \(k\) is very small (e.g., \(k = 1\)) versus very large (e.g., \(k = 1000\)). How does \(k\) affect the relative influence of the top-ranked result versus lower-ranked results?
2. (Coding) Implement a third retriever that uses ripgrep (via subprocess.run(["rg", ...])) for exact string matching and integrate it into the HybridRetriever as a third input to RRF. Test it on queries containing specific error messages or function names. Does adding the ripgrep retriever improve Recall@10 compared to the two-retriever baseline?
3. (Analysis) The U-shape position optimization in Listing 11.18 assumes the "lost in the middle" effect. Design an experiment to measure whether this effect holds for your model of choice: create a context window with a known answer placed at positions 25%, 50%, and 75% through the context, and measure the model's accuracy at each position. Report the results and compare to the findings in Liu et al. (2024).