Skip to content

Build Adaptive RAG From Scratch

Naive RAG treats every query the same way. The pre-retrieval and post-retrieval pages each measured their techniques individually and found the same pattern over and over: a technique that fixes one specific failure mode does nothing — or actively hurts — on a query that doesn't have that failure mode. This tutorial builds the natural next step: an LLM router that looks at each query first, decides what (if anything) is wrong with it, and applies only the matching fix before retrieval — then reranks the result.

The code lives in rag-cookbook/adaptive-rag/; every block below is pulled live from those exact files at build time.

Read these first

This tutorial assumes you've read Advanced RAG — Pre-retrieval and Advanced RAG — Post-retrieval — it reuses their techniques and their measured scenarios directly, and doesn't re-explain how each one works internally.

What you'll build

flowchart LR
    Q[Query] --> R{LLM Router}
    R -->|CLEAR| N[No transformation]
    R -->|VAGUE| W[Rewrite]
    R -->|COMPOUND| D[Decompose]
    R -->|BROAD| M[Multi-query + RRF]
    R -->|NEEDS_CONTEXT| S[Step-back]
    N --> P[(Candidate pool)]
    W --> P
    D --> P
    M --> P
    S --> P
    P --> X[Cross-encoder rerank]
    X --> A[Final top-k]

Five categories, each mapped to a technique already measured on its own: a query that's already clear gets no transformation at all (the honest majority case); a vague, colloquial query gets rewritten; a compound question gets decomposed; a broad, multi-faceted question gets multi-query + RRF; a narrow parameter value gets step-back's broader context added alongside it. Every candidate pool then goes through the same cross-encoder reranking stage before the final top-k comes out.

1. Install

git clone https://github.com/DhruvMakwana/rag-cookbook.git
cd rag-cookbook/adaptive-rag
python3 -m venv .venv && source .venv/bin/activate
pip install -r requirements.txt

# Core
pypdf>=4.0
numpy>=1.26
requests>=2.31
python-dotenv>=1.0
langchain-text-splitters>=0.3

# Embeddings + cross-encoder reranking — both fully local, no API key
sentence-transformers>=3.0

# LLM providers — the router and every pre-retrieval technique needs
# one. Install whichever you plan to use (see .env.example). Ollama
# needs no SDK.
anthropic>=0.34
openai>=1.40

2. Configure your LLM

cp .env.example .env

Every stage of this pipeline that isn't pure vector math makes an LLM call — the router included — so pick a provider and fill in .env the same way as every other recipe on this site.

3. The router

import re

import llm

ROUTER_PROMPT = """Classify the query below into EXACTLY ONE category, \
based on what problem (if any) it poses for retrieval against a \
technical document. Categories, each with an example:

CLEAR - a single specific fact that is a NAMED thing (an algorithm \
name, a hardware type, a dataset name) — not a numeric parameter value \
— already phrased using the kind of technical vocabulary the document \
itself would use.
  Example: "What optimizer did they use for training?"

VAGUE - a single specific fact whose complete answer is ONE short \
thing, but the question is phrased colloquially, avoiding the \
technical vocabulary the document actually uses.
  Example: "why dont they just use the old-school approach everyone \
used before"

COMPOUND - explicitly asks for two or more distinct, independently \
answerable facts, usually joined by "and".
  Example: "What optimizer did they use, and how many layers does it have?"

BROAD - one open-ended question whose COMPLETE answer requires citing \
several distinct reasons or mechanisms (not one short fact) — decide \
based on whether a full answer needs multiple points, regardless of \
whether the wording sounds casual or technical.
  Example: "What lets this model handle relationships between things \
that are positioned far apart, without losing track over long spans?"

NEEDS_CONTEXT - asks for one narrow NUMERIC or configuration parameter \
value (a rate, a dimension, a count) where knowing the broader \
principle/section it belongs to would help interpret or locate the \
answer.
  Example: "What value did they use for the regularization strength?"

Query: {query}

Respond with ONLY the category name, wrapped in <category></category> tags."""


def route_query(question: str, provider: str | None = None) -> str:
    """One LLM call classifies the query so the pipeline can apply only
    the pre-retrieval fix that category actually needs, instead of
    running every technique on every query."""
    response = llm.generate(ROUTER_PROMPT.format(query=question), provider=provider)
    match = re.search(r"<category>(.*?)</category>", response, re.S)
    category = match.group(1).strip() if match else "CLEAR"
    return category if category in {"CLEAR", "VAGUE", "COMPOUND", "BROAD", "NEEDS_CONTEXT"} else "CLEAR"

One LLM call classifies the query into exactly one of 5 categories, each with a one-line example baked directly into the prompt. That's not decoration — an earlier version of this prompt used definitions alone (no examples), and it confused CLEAR with NEEDS_CONTEXT and VAGUE with BROAD on the very queries meant to demonstrate each category. Adding one concrete example per category fixed all of it: 5/5 correct, stable across repeated runs, and correctly generalizing to novel queries never shown in the prompt (tested separately — a colloquially-phrased "why not just make it bigger" question still correctly routes to VAGUE, a genuinely open-ended "why is this more interpretable" question still routes to BROAD).

The category comes back wrapped in <category> tags — the same pattern used for listwise reranking: let the model reason if it wants, give it one unambiguous place to put the actual answer.

4. The full pipeline

import re

import numpy as np
from sentence_transformers import CrossEncoder, SentenceTransformer

import llm

ROUTER_PROMPT = """Classify the query below into EXACTLY ONE category, \
based on what problem (if any) it poses for retrieval against a \
technical document. Categories, each with an example:

CLEAR - a single specific fact that is a NAMED thing (an algorithm \
name, a hardware type, a dataset name) — not a numeric parameter value \
— already phrased using the kind of technical vocabulary the document \
itself would use.
  Example: "What optimizer did they use for training?"

VAGUE - a single specific fact whose complete answer is ONE short \
thing, but the question is phrased colloquially, avoiding the \
technical vocabulary the document actually uses.
  Example: "why dont they just use the old-school approach everyone \
used before"

COMPOUND - explicitly asks for two or more distinct, independently \
answerable facts, usually joined by "and".
  Example: "What optimizer did they use, and how many layers does it have?"

BROAD - one open-ended question whose COMPLETE answer requires citing \
several distinct reasons or mechanisms (not one short fact) — decide \
based on whether a full answer needs multiple points, regardless of \
whether the wording sounds casual or technical.
  Example: "What lets this model handle relationships between things \
that are positioned far apart, without losing track over long spans?"

NEEDS_CONTEXT - asks for one narrow NUMERIC or configuration parameter \
value (a rate, a dimension, a count) where knowing the broader \
principle/section it belongs to would help interpret or locate the \
answer.
  Example: "What value did they use for the regularization strength?"

Query: {query}

Respond with ONLY the category name, wrapped in <category></category> tags."""


def route_query(question: str, provider: str | None = None) -> str:
    response = llm.generate(ROUTER_PROMPT.format(query=question), provider=provider)
    match = re.search(r"<category>(.*?)</category>", response, re.S)
    category = match.group(1).strip() if match else "CLEAR"
    return category if category in {"CLEAR", "VAGUE", "COMPOUND", "BROAD", "NEEDS_CONTEXT"} else "CLEAR"


REWRITE_PROMPT = """Rewrite this query to be clearer and more specific \
for searching a technical document. Return ONLY the rewritten query, \
nothing else.

Query: {query}"""

MULTI_QUERY_PROMPT = """Generate 3 different phrasings of the question \
below, each capturing a different angle on the same underlying \
information need. Return one per line, nothing else.

Question: {query}"""

STEP_BACK_PROMPT = """What is a more general question that would help \
answer this specific question? Return ONLY the general question.

Specific question: {query}"""

DECOMPOSE_PROMPT = """If the question below is asking multiple distinct \
things, break it into separate sub-questions, one per line. If it's \
already a single question, return it unchanged on one line.

Question: {query}"""

CANDIDATE_POOL = 15
FINAL_K = 3
CROSS_ENCODER_MODEL_NAME = "cross-encoder/ms-marco-MiniLM-L-6-v2"


def _bi_encoder_retrieve(question, chunks, embed_model, k):
    query_vec = embed_model.encode(question, convert_to_numpy=True)
    chunk_vecs = embed_model.encode(chunks, convert_to_numpy=True, show_progress_bar=False)
    sims = [float(np.dot(query_vec, v) / (np.linalg.norm(query_vec) * np.linalg.norm(v) + 1e-8)) for v in chunk_vecs]
    top_indices = np.argsort(sims)[::-1][:k]
    return [chunks[i] for i in top_indices]


def _reciprocal_rank_fusion(ranked_lists, k=60):
    scores: dict[str, float] = {}
    for ranked in ranked_lists:
        for rank, chunk in enumerate(ranked):
            scores[chunk] = scores.get(chunk, 0.0) + 1.0 / (k + rank + 1)
    return sorted(scores, key=lambda c: -scores[c])


def _cross_encoder_rerank(reranker, question, candidates, top_n):
    scores = reranker.predict([(question, c) for c in candidates])
    ranked = [c for _, c in sorted(zip(scores, candidates), key=lambda x: -x[0])]
    return ranked[:top_n]


def adaptive_retrieve(
    question: str,
    chunks: list[str],
    embed_model: SentenceTransformer,
    provider: str | None = None,
) -> dict:
    """Router -> matching pre-retrieval technique (or none) -> wide
    candidate pool -> cross-encoder rerank.

    Two design decisions that only surfaced from testing the FULL
    pipeline end to end, not each technique in isolation:

    1. A cross-encoder suffers the SAME vocabulary-mismatch problem a
       bi-encoder does — for VAGUE queries, reranking must be scored
       against the REWRITTEN query, not the original colloquial one, or
       it undoes the fix rewriting just made.
    2. Reranking a decomposed COMPOUND question's merged candidate pool
       against one single query collapses it back to one topic —
       undoing decomposition's whole guarantee of independent
       per-sub-question representation. Each sub-question needs its
       own rerank pass and its own guaranteed slots."""
    reranker = CrossEncoder(CROSS_ENCODER_MODEL_NAME)
    category = route_query(question, provider=provider)

    if category == "VAGUE":
        rewritten = llm.generate(REWRITE_PROMPT.format(query=question), provider=provider).strip()
        candidates = _bi_encoder_retrieve(rewritten, chunks, embed_model, CANDIDATE_POOL)
        final_chunks = _cross_encoder_rerank(reranker, rewritten, candidates, FINAL_K)

    elif category == "COMPOUND":
        response = llm.generate(DECOMPOSE_PROMPT.format(query=question), provider=provider)
        sub_questions = [q.strip("- ").strip() for q in response.split("\n") if q.strip()]
        per_question_pool = max(3, CANDIDATE_POOL // len(sub_questions))
        per_question_final = max(2, FINAL_K // len(sub_questions))
        final_chunks, seen = [], set()
        for sub_q in sub_questions:
            pool = _bi_encoder_retrieve(sub_q, chunks, embed_model, per_question_pool)
            for chunk in _cross_encoder_rerank(reranker, sub_q, pool, per_question_final):
                if chunk not in seen:
                    final_chunks.append(chunk)
                    seen.add(chunk)

    elif category == "BROAD":
        response = llm.generate(MULTI_QUERY_PROMPT.format(query=question), provider=provider)
        variations = [q.strip("- ").strip() for q in response.split("\n") if q.strip()]
        variations.append(question)
        ranked_lists = [_bi_encoder_retrieve(v, chunks, embed_model, CANDIDATE_POOL) for v in variations]
        candidates = _reciprocal_rank_fusion(ranked_lists)[:CANDIDATE_POOL]
        final_chunks = _cross_encoder_rerank(reranker, question, candidates, FINAL_K)

    elif category == "NEEDS_CONTEXT":
        general_question = llm.generate(STEP_BACK_PROMPT.format(query=question), provider=provider).strip()
        specific_hits = _bi_encoder_retrieve(question, chunks, embed_model, CANDIDATE_POOL)
        general_hits = _bi_encoder_retrieve(general_question, chunks, embed_model, CANDIDATE_POOL)
        combined, seen = [], set()
        for chunk in specific_hits + general_hits:
            if chunk not in seen:
                combined.append(chunk)
                seen.add(chunk)
        final_chunks = _cross_encoder_rerank(reranker, question, combined[:CANDIDATE_POOL], FINAL_K)

    else:  # CLEAR
        candidates = _bi_encoder_retrieve(question, chunks, embed_model, CANDIDATE_POOL)
        final_chunks = _cross_encoder_rerank(reranker, question, candidates, FINAL_K)

    return {"category": category, "chunks": final_chunks}

Route, apply the matching pre-retrieval technique, rerank. Two design decisions in this function only became visible once the whole pipeline was tested end-to-end — neither the pre-retrieval page nor the post-retrieval page had any reason to surface them, because each tested its own technique's own top-k directly, with no second stage stacked on top:

A cross-encoder has the same vocabulary-mismatch problem a bi-encoder does. For a VAGUE query, if the final reranking step scores candidates against the original colloquial wording instead of the rewritten one, it throws away the exact fix rewriting just made — the reranker reads the raw query, doesn't recognize the technical passage as a match, and can rank it out of the final top-k even though it's sitting right there in the candidate pool. The fix: VAGUE reranks against the rewritten query, not the original.

Reranking a decomposed question's merged pool against one query undoes decomposition's whole point. COMPOUND retrieval gathers a candidate pool per sub-question specifically so neither fact crowds out the other. Rerank that merged pool with a single cross-encoder pass scored against the full compound question, and it collapses right back down to whichever sub-topic the reranker happens to think is more relevant — silently dropping the other fact with no error. The fix: COMPOUND reranks per sub-question, with a guaranteed minimum number of slots for each one, before combining.

Run it

python adaptive_rag.py --query "What optimizer did they use for training?"
python adaptive_rag.py --compare

--compare runs 5 queries — one per router category, each reusing a scenario already keyword-verified on the pre-retrieval page — through both naive retrieval and the full adaptive pipeline, and prints a three-column comparison: whether naive found the answer, whether the pipeline's pre-retrieval step got it into the candidate pool, and whether it survived all the way to the final output.

Actual output:

Question                                                Category       Naive   Pool    Final
----------------------------------------------------------------------------------------------
What optimizer did they use for training?               CLEAR          True    True    True
why dont they just use RNNs like everyone else did...   VAGUE          False   True    False
What optimizer did they use, and how many attentio...   COMPOUND       False   True    True
What lets this model connect words that are far ap...   BROAD          False   True    False
What value did they use for Pdrop during training?      NEEDS_CONTEXT  True    True    True

What this table actually shows

Three of five queries work end-to-end: CLEAR and NEEDS_CONTEXT correctly get routed to little-or-no transformation and tie naive retrieval (no harm on queries that didn't need fixing); COMPOUND is a clean win, with the per-sub-question reranking fix from above doing real work.

The other two are the more interesting result. VAGUE and BROAD both show Pool = True, Final = False — the router correctly diagnosed the query, correctly applied the fix, and the fix genuinely worked: the target passage entered the candidate pool, something naive retrieval never achieved on either query. And then cross-encoder reranking, scored fairly against the right query, still didn't promote it into the final top-3. This isn't a bug being papered over — it was checked directly: widening the final budget from top-3 to top-5 didn't recover it either. The cross-encoder simply disagrees with what's actually the most precise answer, favoring more general, topically-adjacent passages over the one with the exact technical justification.

That's the honest lesson this whole tutorial builds toward: stacking techniques is not automatically monotonic. Fixing recall at one stage doesn't guarantee precision at the next one — every stage in a pipeline needs its own verification, not just each individual technique tested in isolation the way the earlier pages did. A pipeline that looks correct stage-by-stage can still lose a fix between stages, silently, with no error thrown anywhere.

Adding contextual compression

Compression isn't wired into the default pipeline — it's a separate cost/quality trade-off that not every query needs (same reasoning as the post-retrieval page). Add it as a final step with a flag:

python adaptive_rag.py --query "What optimizer did they use for training?" --compress

Where this still breaks

  • The router itself can misclassify. It's one LLM call with no verification step — a genuinely ambiguous query could land in the wrong category, applying the wrong fix (or none) to a query that needed one.
  • Reranking's final-stage loss on VAGUE/BROAD is unresolved, not just undocumented. This tutorial reports it honestly rather than tuning parameters until it disappears — a real production system would need to either accept a wider final-k for transformed queries, try a different reranker, or measure whether this specific tension shows up on its own corpus at all.
  • This pipeline never checks whether retrieval actually found anything usable before generating. That's a different failure mode entirely — see Self-RAG and Corrective RAG for retrieval-quality checks that happen after this whole pipeline runs, not before.

Sources & further reading