方寸 Portal

AI Engineer Portal

Your personal operating system for career transition.

Private mode
Back to Learning Paths
Lesson

Retrieval patterns

Implement semantic search, hybrid BM25 plus vector retrieval, reranking, and diversity strategies to maximize context quality.

45 min
rag-systemsphase-1portfolio
Prerequisites

Retrieval patterns

Why this matters

Pure semantic search fails in predictable ways. If a user asks about "GPT-4 pricing changes in March 2024," a vector search will retrieve semantically similar content about LLM pricing in general — but might miss the specific document that contains the exact date and model name, because exact keyword match is not what embedding similarity optimizes for.

Production RAG systems use layered retrieval strategies. The goal is not to pick the best single retrieval method; it is to combine methods in a way that finds the right chunks reliably across the full range of query types your users will send. This lesson covers the practical patterns that separate toy demos from production systems.

Core concepts

Semantic search. Query and document embeddings are compared by cosine similarity. Best for natural language questions where the user does not know the exact terminology in the source document. Weak for queries with specific named entities, model numbers, dates, or code identifiers.

BM25 (keyword search). Classic TF-IDF-based ranking that scores documents by exact term overlap, adjusting for document length and term frequency. BM25 does not understand meaning but it is reliable for specific terms. It will find a document containing "GPT-4o-mini rate limit" when that exact phrase is in the text.

Hybrid search. Combine BM25 and vector scores using Reciprocal Rank Fusion (RRF) or a weighted blend. RRF is the most practical approach: each method produces a ranked list, and the final score for each document is the sum of 1/(k + rank) across methods (k=60 is a common default). This handles diverse query types without needing to tune weights per query.

Reranking. A cross-encoder model takes each (query, chunk) pair and scores them jointly. This is more expensive than embedding similarity (which uses independent encodings) but much more accurate. Cohere's Rerank API and cross-encoder models from sentence-transformers are common choices. Typical pattern: retrieve the top 20 chunks cheaply, rerank to find the best 5.

Maximum Marginal Relevance (MMR). When multiple retrieved chunks say the same thing (near-duplicate passages), your context window fills with redundant information. MMR adds a diversity penalty to reduce repetition: each new chunk is selected based on its relevance to the query minus its similarity to already-selected chunks. Use MMR when your source documents are repetitive (multiple versions of the same content, FAQs with overlapping answers).

Metadata filtering. Pre-filter by metadata before semantic search. A query like "find contract terms from agreements signed after 2023" should filter year > 2023 first, then do semantic search within that subset. Filtering reduces noise and speeds up retrieval. This requires good metadata at ingestion time.

Working example

from __future__ import annotations
import math
from collections import defaultdict
from dataclasses import dataclass, field
from typing import Callable

from openai import OpenAI
import chromadb


@dataclass
class RetrievedChunk:
    id: str
    text: str
    metadata: dict
    score: float


def reciprocal_rank_fusion(
    ranked_lists: list[list[RetrievedChunk]],
    k: int = 60,
) -> list[RetrievedChunk]:
    """
    Merge multiple ranked lists using Reciprocal Rank Fusion.
    Returns a single merged list sorted by combined RRF score.
    """
    scores: dict[str, float] = defaultdict(float)
    chunks_by_id: dict[str, RetrievedChunk] = {}

    for ranked_list in ranked_lists:
        for rank, chunk in enumerate(ranked_list):
            scores[chunk.id] += 1.0 / (k + rank + 1)
            chunks_by_id[chunk.id] = chunk

    merged = sorted(chunks_by_id.values(), key=lambda c: scores[c.id], reverse=True)
    # Attach final RRF score for inspection
    for chunk in merged:
        chunk.score = scores[chunk.id]
    return merged


def mmr_rerank(
    query_embedding: list[float],
    candidates: list[RetrievedChunk],
    embed_fn: Callable[[str], list[float]],
    top_k: int = 5,
    diversity_weight: float = 0.3,
) -> list[RetrievedChunk]:
    """
    Maximum Marginal Relevance: balance relevance with diversity.
    diversity_weight: 0.0 = pure relevance, 1.0 = pure diversity.
    """
    def cosine(a: list[float], b: list[float]) -> float:
        dot = sum(x * y for x, y in zip(a, b))
        mag_a = math.sqrt(sum(x**2 for x in a))
        mag_b = math.sqrt(sum(x**2 for x in b))
        return dot / (mag_a * mag_b + 1e-9)

    candidate_embeddings = {c.id: embed_fn(c.text) for c in candidates}
    selected: list[RetrievedChunk] = []
    remaining = list(candidates)

    while len(selected) < top_k and remaining:
        scores = []
        for chunk in remaining:
            relevance = cosine(query_embedding, candidate_embeddings[chunk.id])
            if not selected:
                redundancy = 0.0
            else:
                redundancy = max(
                    cosine(candidate_embeddings[chunk.id], candidate_embeddings[s.id])
                    for s in selected
                )
            mmr_score = (1 - diversity_weight) * relevance - diversity_weight * redundancy
            scores.append((chunk, mmr_score))

        best_chunk, _ = max(scores, key=lambda x: x[1])
        selected.append(best_chunk)
        remaining.remove(best_chunk)

    return selected


class HybridRetriever:
    """
    Combines vector (semantic) and keyword (BM25-approximate) retrieval
    using Reciprocal Rank Fusion, with optional reranking.
    """

    def __init__(self, collection: chromadb.Collection, openai_client: OpenAI):
        self.collection = collection
        self.openai = openai_client
        # Simple in-memory BM25 approximation via Chroma full-text (real
        # production would use Elasticsearch or a dedicated BM25 lib)
        self._texts: list[dict] = []

    def add_chunks(self, chunks: list[dict]) -> None:
        texts = [c["text"] for c in chunks]
        resp = self.openai.embeddings.create(
            model="text-embedding-3-small", input=texts, dimensions=256
        )
        self.collection.add(
            ids=[c["id"] for c in chunks],
            embeddings=[item.embedding for item in resp.data],
            documents=texts,
            metadatas=[c.get("metadata", {}) for c in chunks],
        )
        self._texts.extend(chunks)

    def _vector_search(self, query: str, top_k: int = 20) -> list[RetrievedChunk]:
        resp = self.openai.embeddings.create(
            model="text-embedding-3-small", input=query, dimensions=256
        )
        results = self.collection.query(
            query_embeddings=[resp.data[0].embedding], n_results=top_k
        )
        return [
            RetrievedChunk(id=id_, text=text, metadata=meta, score=1 - dist)
            for id_, text, meta, dist in zip(
                results["ids"][0],
                results["documents"][0],
                results["metadatas"][0],
                results["distances"][0],
            )
        ]

    def _keyword_search(self, query: str, top_k: int = 20) -> list[RetrievedChunk]:
        """Approximate keyword search: score by term overlap with query words."""
        query_terms = set(query.lower().split())
        scored = []
        for chunk in self._texts:
            chunk_terms = set(chunk["text"].lower().split())
            overlap = len(query_terms & chunk_terms)
            if overlap > 0:
                scored.append((chunk, overlap))
        scored.sort(key=lambda x: x[1], reverse=True)
        return [
            RetrievedChunk(id=c["id"], text=c["text"], metadata=c.get("metadata", {}), score=s)
            for c, s in scored[:top_k]
        ]

    def retrieve(
        self,
        query: str,
        top_k: int = 5,
        where: dict | None = None,
    ) -> list[RetrievedChunk]:
        vector_results = self._vector_search(query, top_k=20)
        keyword_results = self._keyword_search(query, top_k=20)

        # Apply metadata filter after retrieval
        if where:
            def matches(chunk: RetrievedChunk) -> bool:
                return all(chunk.metadata.get(k) == v for k, v in where.items())
            vector_results = [c for c in vector_results if matches(c)]
            keyword_results = [c for c in keyword_results if matches(c)]

        merged = reciprocal_rank_fusion([vector_results, keyword_results])
        return merged[:top_k]

Common mistakes

  1. Only using vector search. For anything with specific entities, codes, dates, or technical identifiers, pure vector search misses. Add keyword search before your launch date, not after users complain.

  2. Not reranking when you have budget. Retrieving top-20 cheaply and reranking to top-5 is standard practice. The quality improvement is significant; the cost is a Cohere Rerank call per query.

  3. Skipping metadata filtering. If users can scope queries ("only show me content from the 2024 handbook"), build metadata filtering into the retrieval layer on day one. Adding it later requires schema migrations.

  4. Retrieving too few candidates. If you retrieve top-3 and one is irrelevant, you only have 2 useful chunks in context. Retrieve top-10 or top-20, filter for quality, then pass the best 3-5 to the model.

  5. Never looking at what was retrieved. Log the chunks returned for a sample of queries. You will immediately see what is broken. Most engineers skip this step and chase quality problems in the wrong layer.

Try it yourself

Build a retrieval harness that logs both the query and every chunk returned, with its distance score. Run 10 queries from your domain. Identify the query types where vector search succeeds and where it fails. Try adding a keyword boost for exact entity matches. Measure how often the chunk that should answer the question appears in the top-3 results versus the top-10 results.

Lesson Deep Dive

Ask a follow-up question about this lesson and get an AI-powered explanation.

Loading previous questions...