AI Engineer Portal

Your personal operating system for career transition.

Private mode

Exercise

Implement query expansion for improved RAG recall

Implement Query Expansion for Improved RAG Recall

Users do not always phrase their questions the way your documents phrase their answers. A user asks "how do I cancel my subscription" but the documentation says "to terminate your membership agreement." Embedding similarity handles common paraphrases well, but for domain-specific or legal/medical terminology gaps, generating multiple query variants and merging results is more reliable.

What you are building

Create a QueryExpander class that:

Generates query variants using an injectable LLM function — given the original query, produce N alternative phrasings with different vocabulary.
Retrieves candidates for each variant (including the original) using an injectable retrieve function.
Merges and deduplicates all results using Reciprocal Rank Fusion (RRF) with k=60.
Returns a QueryExpansionResult with: original query, expanded queries, merged results with RRF scores, per-variant results for debugging, and found_by_queries count per document.

Why this matters

Query expansion consistently improves recall for domain-specific queries. Documents appearing in results for multiple query variants are likely genuinely relevant — multi-variant agreement is a strong relevance signal. The cost is N small-model LLM calls plus N retrieve calls per query.

Constraints

LLM function: Callable[[str, int], list[str]] — query + n -> variants.
Retrieve function: Callable[[str, int], list[tuple[str, float]]] — query + k -> [(doc_id, score)].
Use RRF k=60. Type-hint everything.

Retrieval / medium / Step 12 of 15

Practice stage

Retrieval quality and ranking

Hint

Do not trust a single score blindly. Combine ranking logic with metadata and think about why a candidate deserves to rise.

Next drill

Build a production RAG pipeline with caching and fallbacks

Success criteria

- Ranking rule is explainable
- Supports future iteration without rewriting everything
- Feels tied to product trust rather than pure math trivia

Review checklist

- Did I account for both relevance and metadata?
- Would I be able to explain the ranking rule in an interview?
- Does the code make future tuning easy?

from __future__ import annotations

from collections import defaultdict
from dataclasses import dataclass, field
from typing import Callable

@dataclass
class ExpandedResult:
    doc_id: str
    rrf_score: float
    found_by_queries: list[str]

@dataclass
class QueryExpansionResult:
    original_query: str
    expanded_queries: list[str]
    merged_results: list[ExpandedResult]
    per_variant: dict[str, list[tuple[str, float]]]

def rrf_merge(
    ranked_lists: list[list[tuple[str, float]]],
    k: int = 60,
) -> list[tuple[str, float]]:
    """Merge ranked lists: score(doc) = sum(1 / (k + rank)) across all lists."""
    # TODO: implement
    raise NotImplementedError

class QueryExpander:
    def __init__(
        self,
        expand_fn: Callable[[str, int], list[str]],
        retrieve_fn: Callable[[str, int], list[tuple[str, float]]],
        n_variants: int = 3,
        top_k_per_query: int = 10,
    ):
        self.expand_fn = expand_fn
        self.retrieve_fn = retrieve_fn
        self.n_variants = n_variants
        self.top_k_per_query = top_k_per_query

def retrieve(self, query: str) -> QueryExpansionResult:
        """Expand query, retrieve for all variants, merge results with RRF."""
        # TODO: implement
        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.