Phase 3: RAG & Knowledge Systems

Retrieval, reranking & generation flow

Advanced ~17 min read
Think of it this way A friendly analogy. Read this if the technical version feels dense. Show Hide

Imagine you're making a super special birthday cake, one of those amazing ones with lots of layers and decorations. You wouldn't just throw everything from your kitchen into a bowl and hope for the best, right? Making something awesome, especially with computers, is a lot like that. We break it down into different, important steps, and knowing these steps helps us make sure our creations are really good.

The very first step is like gathering your ingredients. Before you even start baking, you'd go to the store or your pantry and just grab everything you think might be needed for a cake. Flour, sugar, eggs, sprinkles, chocolate, maybe even some weird spices just in case! You want to make sure you don't miss anything important, so you grab a lot, even if some things won't end up in the final cake. This step is all about being fast and making sure you have plenty of options to choose from.

Once you have a big pile of possibilities, the second step is picking the perfect ones. Now you look at your actual recipe and carefully choose. "Okay, the recipe calls for two cups of this specific kind of flour, not that other one. And I need three fresh eggs, not those ones that have been in the fridge too long." You're being really picky here, checking each ingredient to make sure it's exactly what you need and in the best condition. This takes a little more time and thought than just grabbing, but it makes sure your cake will taste amazing because you're using only the very best.

Finally, with your perfect ingredients neatly selected, the third step is making the cake itself. You follow the recipe, mix everything in the right order, and bake it. You don't suddenly decide to add ketchup or invent a new ingredient that wasn't in your chosen pile. You use only the excellent ingredients you carefully picked, following the instructions to create a delicious, finished cake. This whole process, from gathering to baking, means you get a much better result than if you just tried to do everything all at once. So, when you're building smart computer programs later on, remembering these three stages will help you make sure your creations are really top-notch!

The mental model: three pipelines stitched together

Think of RAG as three sub-systems with different optimization targets. The retriever is optimized for recall: it needs to surface every chunk that could plausibly answer the question. The reranker is optimized for precision: it reads the query and each candidate together (not independently, as the bi-encoder retriever does) and assigns a calibrated relevance score. The generator is optimized for faithfulness: it should produce a response that is directly supported by the context it receives, not by parametric memory. Each sub-system has its own latency, cost, and accuracy profile. Conflating them leads to bad tuning decisions. For example, cutting TOP_K_RETRIEVE from 20 to 5 to "save time" actually hurts the reranker because it has fewer candidates to pick from, and the relevant chunk may have ranked 8th in the first stage.

The retrieval stage in detail

Vector search with an ANN index (HNSW is the common choice in Qdrant, Pinecone, and Weaviate) is a bi-encoder architecture: your query is encoded independently, your stored chunks were encoded independently at ingest time, and you compute approximate cosine similarity between them. This is fast, O(log n) at query time, and scales to tens of millions of vectors. The limitation is that the model never sees the query and chunk together, so it can miss lexical cues. A common fix is hybrid search: combine dense ANN results with sparse BM25 keyword results using Reciprocal Rank Fusion (RRF). Qdrant's built-in sparse vectors and Weaviate's BM25 integration both support this. If your documents are terminology-heavy (legal, medical, code), hybrid search will improve recall measurably compared to dense-only.

The reranking stage in detail

A cross-encoder reads (query, chunk) as a single concatenated input and produces one relevance score. This is far more accurate than cosine similarity because the model can attend across both texts simultaneously. The tradeoff is that you cannot pre-compute these scores at ingest time. Every query requires a fresh forward pass over all TOP_K_RETRIEVE candidates, so cross-encoders are much more expensive per query. In practice, running a small cross-encoder like cross-encoder/ms-marco-MiniLM-L-6-v2 on 20 candidates takes roughly 50-100 ms on CPU. Cohere's Rerank API (cohere.rerank()) offloads this to their infrastructure and consistently outperforms open-source cross-encoders on general domains. For domains where you have labeled data, fine-tuning a cross-encoder on (query, positive chunk, negative chunk) triplets is the highest-leverage accuracy improvement available in a RAG pipeline.

The generation stage in detail

Your top-k reranked chunks get packed into the system prompt or a dedicated context block. This is where token budget management becomes critical. If each chunk is 400 tokens and you pass 10 chunks, you've consumed 4,000 tokens of context before the user even writes a word. That's fine for a 128k context model, but it directly impacts cost and latency. A production heuristic: target 3-6 chunks, each under 512 tokens, for a total context cost under 3k tokens. Beyond accuracy, you also need to decide on the context structure. Interleaving chunk source metadata (document title, page number) inside the context block lets the model cite sources. A system prompt instruction like "Always cite [Source: <title>] after each claim" dramatically improves traceability. The alternative, stripping metadata and hoping the model summarizes cleanly, makes debugging impossible when the answer is wrong.

Real-world scenario: an internal knowledge base for a SaaS support team

You're building a support copilot that answers questions using 5,000 help articles. At 10 users, none of this matters much. At 1,000 concurrent queries per hour, you're now spending real money on Cohere Rerank API calls (each call touches ~20 documents). The practical move is to cache rerank results keyed by (query_hash, top_k_retrieve_result_hashes). Many support questions are semantically identical. A Redis cache with a 24-hour TTL can absorb 60-80% of your reranker traffic if your user base asks predictable questions. At 10M users, you're thinking about dedicated GPU instances for a self-hosted cross-encoder, or batching rerank requests to reduce API overhead. The retrieval stage itself should already be fast from your ANN index, but you may need to shard your vector collection across multiple Qdrant nodes if you're at 50M+ vectors.

Tradeoffs versus alternative approaches

The three-stage pipeline (retrieve -> rerank -> generate) is not the only architecture. A simpler two-stage pipeline (retrieve -> generate, skip reranking) is valid when your TOP_K_RETRIEVE is already small (3-5) and your embedding model is high quality. Skip the reranker when latency is the primary constraint and your documents are relatively uniform in topic. Conversely, if your corpus is noisy or cross-domain, skipping reranking is the most common cause of "the LLM gave a confident but wrong answer" bugs. A newer approach is late interaction models like ColBERT, which store per-token embeddings and do a cheap MaxSim operation at query time. This sits between bi-encoder speed and cross-encoder accuracy. LlamaIndex and RAGatouille both have ColBERT integrations. The downside is index size: ColBERT embeddings are roughly 10x larger than single-vector embeddings.

Key Takeaways

  • Use ANN vector search for high-recall retrieval, then a cross-encoder for high-precision reranking.
  • Pass only reranked, deduplicated chunks to the LLM to avoid context dilution.
  • Log retrieval scores and rerank scores separately so you can diagnose which stage is failing.
  • Token budget for context is finite; always enforce a max-chunk-count before generation.

Pro tips

  • Cross-encoder reranking is the highest ROI accuracy investment in most pipelines, but the gains disappear if your retriever has poor recall. Always verify recall@20 before tuning the reranker. A retriever that misses the right chunk at position 25 cannot be fixed downstream.
  • Log your rerank scores in production. A bimodal distribution (scores clustered near 0 and near 1) means your reranker is confident. A flat distribution between 0.3 and 0.7 means your chunks are all borderline relevant, which usually points to a chunking problem upstream.
  • When the context window is large (e.g., GPT-4o with 128k tokens), the temptation is to stuff 30 chunks in and let the model figure it out. Resist. LLMs still suffer 'lost in the middle' degradation where relevant content in the middle of a long context is underweighted. Fewer, higher-quality chunks consistently outperform more, lower-quality chunks.
  • Hybrid search (dense + sparse BM25) is not always better than dense-only. On short, conversational queries it helps significantly. On long, descriptive queries it can hurt because BM25 over-weights rare terms. A/B test on your actual query distribution before committing to the added infrastructure complexity.

Common pitfalls

  • Mistake: Passing raw retriever scores to the LLM as "relevance percentages." Fix: Retriever cosine scores are not calibrated probabilities. Only expose reranker scores or omit scores entirely.
  • Mistake: Using a single TOP_K value (e.g., 5) for both retrieval and final context. Fix: Retrieve broadly (15-25), rerank, then pass the top 3-6 to the LLM. Separate the two numbers.
  • Mistake: Not deduplicating chunks before generation. If two chunks are 90% identical text from the same doc, you waste tokens and the model may over-weight that source. Fix: Hash chunks at ingest time and deduplicate by hash before packing context.
  • Mistake: Forgetting to include source metadata in the context block. Fix: Always prepend [Source: <title>] to each chunk. Debugging incorrect answers is nearly impossible without knowing which chunk the model cited.

When to use retrieve-only vs retrieve+rerank vs ColBERT

Option Use when Avoid when
Retrieve only (bi-encoder + ANN) Latency is paramount, corpus is small and topically uniform, embedding model quality is high (e.g., text-embedding-3-large). Corpus is cross-domain or noisy, queries are ambiguous, or answer quality is critical and you can afford 50-150 ms extra.
Retrieve + cross-encoder rerank Answer accuracy is the primary metric, corpus is large or noisy, and you can absorb the extra reranking latency and cost. You need sub-100 ms end-to-end response time and cannot afford a second model call or API round trip.
ColBERT late interaction You want accuracy close to a cross-encoder but at bi-encoder query-time speed, and you can absorb a 10x larger index. Your vector store does not support per-token multi-vector indexing, or index storage cost is a hard constraint.
LLM-based reranker (e.g., GPT-4o as judge) Domain is highly specialized, no suitable cross-encoder exists, and query volume is low enough that LLM reranking cost is acceptable. You have more than a few hundred queries per hour. LLM reranking at scale is expensive and slow.

Code Example

python
# openai>=1.0.0, sentence-transformers>=2.7.0, qdrant-client>=1.9.0
import os
from qdrant_client import QdrantClient
from sentence_transformers import SentenceTransformer, CrossEncoder
from openai import OpenAI

DOC_COLLECTION = "kb_chunks"
TOP_K_RETRIEVE = 20
TOP_K_RERANK = 4

retriever = SentenceTransformer("BAAI/bge-small-en-v1.5")
reranker = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2")
qdrant = QdrantClient(url=os.environ["QDRANT_URL"])
client = OpenAI()

def rag(query: str) -> str:
    # Stage 1: Retrieve
    query_vec = retriever.encode(query).tolist()
    hits = qdrant.search(DOC_COLLECTION, query_vec, limit=TOP_K_RETRIEVE)
    texts = [h.payload["text"] for h in hits]

    # Stage 2: Rerank
    pairs = [(query, t) for t in texts]
    scores = reranker.predict(pairs)
    ranked = sorted(zip(scores, texts), reverse=True)[:TOP_K_RERANK]
    context = "\n\n---\n\n".join(t for _, t in ranked)

    # Stage 3: Generate
    resp = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[
            {"role": "system", "content": "Answer using only the context below.\n\n" + context},
            {"role": "user", "content": query},
        ],
    )
    return resp.choices[0].message.content

How this code works

The provided Python code defines a complete Retrieval-Augmented Generation (RAG) pipeline, designed to answer user queries accurately by drawing information from a specific knowledge base. Its job is to ensure an AI model's responses are grounded in provided facts rather than its general training data, enhancing trustworthiness and relevance.

The rag function orchestrates three stages. First, in "Retrieve," the user query is transformed into a numerical query_vec using the retriever (a SentenceTransformer). qdrant.search then uses this vector to efficiently find an initial set of relevant document chunks, limited by TOP_K_RETRIEVE, from the DOC_COLLECTION. Next, in "Rerank," a reranker (a CrossEncoder) takes the original query and each retrieved text chunk. It predicts relevance scores, then sorts and selects the absolute best TOP_K_RERANK chunks. These highly relevant chunks form the context for the final stage. The subtle but critical instruction in the messages for client.chat.completions.create is the "Answer using only the context below" system prompt, which strictly constrains the gpt-4o-mini model to synthesize an answer solely from the provided context, preventing external knowledge or fabrication.

Production-grade example

Adds retries, graceful rerank fallback, streaming generation, per-stage latency logging, and token-budget guard.

python
# openai>=1.0.0, cohere>=5.0.0, qdrant-client>=1.9.0, tenacity>=8.2.0
import os
import time
import logging
import hashlib
import json
from typing import Generator
from tenacity import retry, stop_after_attempt, wait_exponential, retry_if_exception_type
from qdrant_client import QdrantClient
from qdrant_client.http.exceptions import ResponseHandlingException
from openai import OpenAI, APITimeoutError, RateLimitError
import cohere

logger = logging.getLogger(__name__)
qdrant = QdrantClient(url=os.environ["QDRANT_URL"], api_key=os.environ["QDRANT_API_KEY"], timeout=5.0)
co = cohere.Client(os.environ["COHERE_API_KEY"])
oai = OpenAI(api_key=os.environ["OPENAI_API_KEY"])

COLLECTION = "kb_chunks"
RETRIEVE_K = 20
RERANK_K = 5
MODEL = "gpt-4o-mini"
MAX_CONTEXT_TOKENS_ESTIMATE = 3000  # rough; ~4 chars per token

@retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=1, max=8),
       retry=retry_if_exception_type(ResponseHandlingException))
def _retrieve(query_vec: list[float]) -> list[dict]:
    hits = qdrant.search(COLLECTION, query_vector=query_vec, limit=RETRIEVE_K, with_payload=True)
    return [{"text": h.payload["text"], "source": h.payload.get("source", "unknown"),
             "retrieve_score": h.score} for h in hits]

def _rerank(query: str, candidates: list[dict]) -> list[dict]:
    if not candidates:
        return []
    try:
        result = co.rerank(model="rerank-english-v3.0", query=query,
                           documents=[c["text"] for c in candidates], top_n=RERANK_K)
        reranked = []
        for r in result.results:
            doc = candidates[r.index].copy()
            doc["rerank_score"] = r.relevance_score
            reranked.append(doc)
        return reranked
    except cohere.errors.TooManyRequestsError:
        logger.warning("Cohere rate limit hit; falling back to retrieve-order top-%d", RERANK_K)
        return candidates[:RERANK_K]  # graceful degradation

def _build_context(docs: list[dict]) -> str:
    parts = []
    total = 0
    for d in docs:
        chunk = f"[Source: {d['source']}]\n{d['text']}"
        total += len(chunk)
        if total > MAX_CONTEXT_TOKENS_ESTIMATE * 4:
            break
        parts.append(chunk)
    return "\n\n---\n\n".join(parts)

@retry(stop=stop_after_attempt(2), wait=wait_exponential(multiplier=1, min=2, max=10),
       retry=retry_if_exception_type((RateLimitError, APITimeoutError)))
def _generate_stream(query: str, context: str) -> Generator[str, None, None]:
    system = (
        "You are a helpful assistant. Answer the user's question using ONLY the context below. "
        "Cite sources as [Source: <name>] after each claim. "
        "If the context does not contain the answer, say 'I don't have enough information.'\n\n"
        + context
    )
    stream = oai.chat.completions.create(
        model=MODEL, messages=[{"role": "system", "content": system},
                                {"role": "user", "content": query}],
        stream=True, timeout=30.0
    )
    for chunk in stream:
        delta = chunk.choices[0].delta.content
        if delta:
            yield delta

def rag_query(query: str, query_vec: list[float]) -> Generator[str, None, None]:
    t0 = time.perf_counter()
    candidates = _retrieve(query_vec)
    t_retrieve = time.perf_counter()

    reranked = _rerank(query, candidates)
    t_rerank = time.perf_counter()

    context = _build_context(reranked)
    logger.info(json.dumps({"event": "rag_stages", "query_hash": hashlib.md5(query.encode()).hexdigest(),
                            "retrieve_count": len(candidates), "rerank_count": len(reranked),
                            "retrieve_ms": round((t_retrieve - t0) * 1000),
                            "rerank_ms": round((t_rerank - t_retrieve) * 1000),
                            "top_rerank_score": reranked[0]["rerank_score"] if reranked else None}))
    yield from _generate_stream(query, context)

How this code works

This code creates a robust Retrieval-Augmented Generation (RAG) pipeline, designed to answer user questions by leveraging a specific knowledge base. The rag_query function orchestrates the entire process: it takes a user's question and its vector representation, retrieves potentially relevant documents, reranks them to find the most pertinent information, constructs a focused context, and then uses a large language model (LLM) to generate an answer, streaming it back token by token. This structured approach helps ensure the LLM's responses are accurate and grounded in provided data.

The process begins with _retrieve, fetching an initial set of RETRIEVE_K (20) document chunks from Qdrant. These are then passed to _rerank, which employs Cohere's rerank-english-v3.0 model to identify the RERANK_K (5) most relevant documents, significantly improving the quality of the information. A subtle but important detail occurs in _build_context: it meticulously combines the reranked documents, adding their sources, but critically stops adding more content if the estimated total length, using MAX_CONTEXT_TOKENS_ESTIMATE * 4, risks exceeding the language model's input capacity. This proactive step prevents common token limit errors. Finally, _generate_stream sends this curated context and the original query to an OpenAI model, providing a real-time streamed response. Error handling and retries are built into critical API calls using the @retry decorator for resilience.

Practice & master

Try the exercise, check your understanding, then mark this lesson mastered to track your path to pro.

Exercise

Wire together a minimal three-stage RAG function that retrieves 10 chunks from a pre-populated Qdrant collection, reranks them with the cross-encoder/ms-marco-MiniLM-L-6-v2 model locally, and streams the GPT-4o-mini response to stdout. Log the top reranked chunk's score before generating.

python
# sentence-transformers>=2.7.0, qdrant-client>=1.9.0, openai>=1.0.0
import os
from qdrant_client import QdrantClient
from sentence_transformers import SentenceTransformer, CrossEncoder
from openai import OpenAI

COLLECTION = "kb_chunks"  # assumes this exists in your local Qdrant
RETRIEVE_K = 10
RERANK_K = 3

retriever = SentenceTransformer("BAAI/bge-small-en-v1.5")
reranker = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2")
qdrant = QdrantClient(url=os.environ["QDRANT_URL"])
client = OpenAI()

def rag(query: str) -> None:
    # TODO: Stage 1 - encode query, search qdrant for RETRIEVE_K results
    # hits = ...
    # texts = ...

    # TODO: Stage 2 - build (query, text) pairs, predict scores, sort, take top RERANK_K
    # scores = ...
    # ranked = ...
    # print the top chunk's score here

    # TODO: Stage 3 - build context string, call OpenAI with stream=True, print each chunk
    pass

if __name__ == "__main__":
    rag("What is the refund policy for annual subscriptions?")

Quick check

  1. Why do you retrieve more chunks than you ultimately pass to the LLM?

  2. A cross-encoder reranker is slower than a bi-encoder retriever at query time. Why?

  3. Your RAG system returns accurate answers for common questions but fails on technical jargon queries. Which change is most likely to help?

Self-check: Explain why you would set RETRIEVE_K=20 and RERANK_K=4 rather than just RETRIEVE_K=4. Then describe what metric you would check first if your RAG answers are frequently wrong, and which pipeline stage it implicates.