Retrieval, reranking, and generation

Here is the heart of the system. On every question, we embed the query, pull 20 candidate chunks from ChromaDB, rerank them with a cross-encoder, pass the top 5 into the prompt as context, and stream the grounded answer back to the user.

Retrieval and reranking

Two-stage retrieval: fast recall with a bi-encoder, precise reranking with a cross-encoder.

rag_utils.py
python
class Reranker:
    """Reranks retrieved documents to improve relevance"""

    def __init__(self, model_name: str = "cross-encoder/ms-marco-MiniLM-L-6-v2"):
        self.model = CrossEncoder(model_name)

    def rerank(self, query: str, docs, top_k: int = 5):
        if not docs:
            return []
        pairs = [[query, doc.content] for doc in docs]
        scores = self.model.predict(pairs)
        doc_scores = list(zip(docs, scores))
        doc_scores.sort(key=lambda x: x[1], reverse=True)
        return [doc for doc, score in doc_scores[:top_k]]

The cross-encoder compares the query against each candidate directly. It is slower than bi-encoder similarity, but reranking 20 candidates is cheap and the precision gain is substantial.

rag_utils.py
python
class SimpleRAGPipeline:
    def __init__(self, llm_provider, embedding_provider, vector_store):
        self.llm_provider = llm_provider
        self.embedding_provider = embedding_provider
        self.vector_store = vector_store
        self.reranker = Reranker()

    async def retrieve(self, query: str, filters=None):
        # 1. Embed the query
        query_embedding = (await self.embedding_provider.generate_embeddings([query]))[0]

        # 2. Pull many candidates for recall
        candidates = self.vector_store.query(
            query_embedding=query_embedding,
            n_results=20,
            where=filters,
        )

        # 3. Rerank for precision
        return self.reranker.rerank(query, candidates, top_k=5)

Retrieve many, rerank to a few. The bi-encoder is fast but approximate. The cross-encoder is accurate but slow. Combined, they give you both speed and quality.

Vector similarity is a rough heuristic. Two chunks can have high cosine similarity while only one actually answers the question. The cross-encoder compares the exact query against each candidate text, which is a much more accurate signal of relevance. In practice, reranking often surfaces the real answer from rank 6 or 12 up to rank 1.

Matching exercise: Match each component to its job

Loading practice…