AI Engineer Portal

Your personal operating system for career transition.

Private mode

Exercise

Add reranking to retrieval results

Add Reranking to Retrieval Results

First-stage retrieval optimizes for recall: get the right documents into the candidate set. Reranking optimizes for precision: from that candidate set, find the documents most useful for this specific query. A cross-encoder reranker scores each (query, document) pair jointly — seeing both at once for finer-grained relevance signals.

Typical production pattern: retrieve top-20 cheaply, rerank to find the best 5.

What you are building

Create a RerankedRetriever that wraps an existing retriever and adds a reranking step:

Retrieve a larger candidate set (e.g., top-20) from the wrapped retriever.
Rerank candidates by calling a rerank_fn(query, documents) -> list[float].
Return the top-k documents by rerank score with original retrieval rank and rerank score attached.
Log rank changes — compute average absolute position change to evaluate whether reranking is helping.
Handle reranker failures gracefully — if the rerank function raises, fall back to original order and log the error.

Constraints

The rerank_fn is injected — works with Cohere API, local cross-encoder, or mock.
All methods type-hinted.
RankedResult objects must include original_rank and rerank_score.
Fallback returns original results with rerank_failed: True on the response.

Rag Systems / medium / Step 3 of 6

Practice stage

General drill

Hint

Keep the solution explicit and reviewable.

Next drill

Evaluate RAG faithfulness with citation tracking

Success criteria

Make the solution explicit, debuggable, and easy to explain.

Review checklist

Review where the boundary is, what gets validated, and what would be hard to debug later.

from __future__ import annotations

import logging
from dataclasses import dataclass, field
from typing import Callable, Optional, Protocol

logger = logging.getLogger(__name__)

@dataclass
class Document:
    id: str
    text: str
    metadata: dict = field(default_factory=dict)

@dataclass
class RankedResult:
    document: Document
    original_rank: int
    rerank_score: float
    final_rank: int

@dataclass
class RetrievalResponse:
    results: list[RankedResult]
    rerank_failed: bool = False
    rank_shift_avg: float = 0.0

class FirstStageRetriever(Protocol):
    def retrieve(self, query: str, top_k: int) -> list[Document]: ...

class RerankedRetriever:
    def __init__(
        self,
        base_retriever: FirstStageRetriever,
        rerank_fn: Callable[[str, list[str]], list[float]],
        candidate_k: int = 20,
    ):
        self.base_retriever = base_retriever
        self.rerank_fn = rerank_fn
        self.candidate_k = candidate_k

def retrieve(self, query: str, top_k: int = 5) -> RetrievalResponse:
        """Two-stage retrieve + rerank."""
        # TODO: retrieve candidate_k candidates from base_retriever
        # TODO: call rerank_fn(query, [doc.text for doc in candidates])
        # TODO: if rerank_fn raises, log and return original results with rerank_failed=True
        # TODO: sort by rerank score descending, take top_k
        # TODO: compute rank_shift_avg: mean of |original_rank - final_rank|
        raise NotImplementedError

Practice

Generate a variation

Generate a new exercise variation to deepen understanding or practice a related concept.

Attempt history

Recent submissions

Before you submit, decide what a strong answer should make obvious to the reviewer.

No attempts yet.