Capsule 03: Query expansion — when one query isn't enough

Capsule overview

In capsule 02 you saw that user queries are frequently ambiguous, incomplete or badly phrased. The most direct technique for solving the ambiguity problem is query expansion: instead of searching with the original query (a single one, possibly misinterpreted), you generate several versions that cover the likely interpretations, search with each one, and combine the rankings.

The key pedagogical insight: if the user writes "fastapi auth", you don't know whether they want OAuth2, JWT, API keys, basic auth or sessions. But you can generate 5 queries that cover all 5 interpretations, search with each one, and hand the LLM the best results from all 5 searches combined. The probability of capturing the right answer rises from ~50% (with the ambiguous original query) to ~85% (with the 5 expansions).

This capsule teaches you how to generate quality expansions with an LLM (it isn't trivial — a badly designed prompt produces noisy expansions), how to combine the results with Reciprocal Rank Fusion, and how to decide when to expand vs when to use the query directly.

By the end of this capsule you'll be able to:

• ✅ Generate query expansions with an LLM using calibrated prompts
• ✅ Implement Reciprocal Rank Fusion (RRF) to combine multiple rankings
• ✅ Decide when to expand (ambiguous queries) vs when not to (specific queries)
• ✅ Calculate the extra cost (latency + LLM calls + retrieval calls) and justify it
• ✅ Anticipate the most expensive trap: expansions that change the user's intent
• ✅ Implement expansion with a cache for repeated queries

Estimated time: 30-35 minutes

The insight: one query, several interpretations, several searches

Think about how you google when you can't find something. The first query doesn't work, so you try variants:

Query 1: "fastapi auth"             → didn't find what you wanted
Query 2: "fastapi oauth2"           → found something, but not specific
Query 3: "fastapi password bearer"  → found exactly what you wanted

What you do manually with 3 iterative queries, query expansion does automatically with 5 queries in parallel, then combines the results:

       ┌─────────────────────────────────────────────────┐
       │ User query: "fastapi auth"                       │
       └────────────────────┬────────────────────────────┘
                            │ LLM expansion
                            ▼
       ┌─────────────────────────────────────────────────┐
       │ Query 1: "How to implement OAuth2 in FastAPI?"  │
       │ Query 2: "FastAPI JWT authentication"            │
       │ Query 3: "FastAPI API key auth"                  │
       │ Query 4: "FastAPI session-based auth"            │
       │ Query 5: "FastAPI security best practices"       │
       └────────────────────┬────────────────────────────┘
                            │ Each query → search
                            ▼
       ┌──────────────────────────────────────────────────┐
       │ 5 result sets (top-10 each = 50 docs)            │
       └────────────────────┬─────────────────────────────┘
                            │ Reciprocal Rank Fusion
                            ▼
       ┌─────────────────────────────────────────────────┐
       │ The final top-5: docs that rank high across     │
       │ MULTIPLE queries (a strong relevance signal)    │
       └─────────────────────────────────────────────────┘

Why it works: a document that shows up at position #2 for the query "FastAPI OAuth2" and at position #1 for "FastAPI security best practices" carries a stronger relevance signal than one that appears in only one of the searches. RRF captures exactly that intersection of rankings.

Generating quality expansions with an LLM

The prompt is the sensitive part. A badly designed prompt generates expansions that are:

• Too similar (they add no diversity of interpretation)
• Too different (they change the user's intent)
• Too generic (they never use the domain's specific vocabulary)

# query_expansion.py
from openai import OpenAI
from pydantic import BaseModel, Field
from typing import List
import os


client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))


class ExpandedQueries(BaseModel):
    queries: List[str] = Field(
        description="Diverse alternative queries that cover different interpretations of the original",
        min_items=3,
        max_items=8,
    )
    reasoning: str = Field(
        description="One-sentence explanation of how the queries cover different interpretations"
    )


SYSTEM_PROMPT = """You are an expert at expanding search queries for retrieval systems.

Given a short or ambiguous user query, generate 5 alternative queries that:

1. Cover the most likely DIFFERENT interpretations of the original
2. Use natural language (full questions, not keywords)
3. Preserve the user's INTENT — don't add unrelated topics
4. Use specific terminology when relevant (e.g., if the topic is FastAPI, use FastAPI-specific terms)
5. Are mutually distinct (each one explores a different angle)

Examples:

Input: "fastapi auth"
Output queries:
1. How to implement OAuth2 authentication in FastAPI?
2. FastAPI JWT token authentication tutorial
3. How to use API keys for authentication in FastAPI?
4. FastAPI session-based authentication with cookies
5. Best practices for securing FastAPI endpoints

Input: "ml deployment"
Output queries:
1. How to deploy machine learning models to production?
2. ML model deployment with Docker containers
3. Serverless deployment for machine learning models (AWS Lambda, Cloud Run)
4. Real-time vs batch ML model serving
5. CI/CD pipelines for ML model deployment"""


def expand_query(query: str, num_expansions: int = 5) -> ExpandedQueries:
    """Generate diverse query expansions preserving original intent."""
    response = client.beta.chat.completions.parse(
        model="gpt-4o-mini",
        messages=[
            {"role": "system", "content": SYSTEM_PROMPT},
            {
                "role": "user",
                "content": f'User query: "{query}"\n\nGenerate {num_expansions} alternative queries.',
            },
        ],
        response_format=ExpandedQueries,
        temperature=0.5,  # a little creativity, for diversity
    )
    return response.choices[0].message.parsed


# Try it
expanded = expand_query("fastapi auth", num_expansions=5)
print(f"Reasoning: {expanded.reasoning}\n")
for i, q in enumerate(expanded.queries, 1):
    print(f"{i}. {q}")

Expected output:

Reasoning: The original query 'fastapi auth' is ambiguous and could refer to several authentication methods. The expansions cover the main implementation patterns.

1. How to implement OAuth2 authentication in FastAPI applications?
2. FastAPI JWT token authentication implementation guide
3. How to use API key authentication in FastAPI?
4. Implementing session-based authentication in FastAPI with cookies
5. FastAPI security best practices and common authentication patterns

Why structured outputs

Without structured outputs, LLMs tend to:

• Return inconsistent formats (sometimes numbered, sometimes bulleted, sometimes free text)
• Add comments or explanations in the middle that break your parsing
• Return more or fewer queries than you asked for

response_format=ExpandedQueries (with Pydantic) guarantees a clean, parseable list.

Why temperature=0.5

• temperature=0 produces expansions nearly identical to the original query (little diversity).
• temperature=1.0 produces creative expansions that sometimes wander off-topic.
• 0.5 is the sweet spot: diverse, but they respect the intent.

Reciprocal Rank Fusion (RRF): the right way to fuse

After searching with the 5 queries, you have 5 result sets. How do you combine 5 rankings into one?

Bad option 1 — sum the scores: cosine scores from different queries aren't comparable. Adding them gives you meaningless rankings.

Bad option 2 — use only the top-1 from each: it throws away every other result. If a document shows up as #2 in 4 different queries, that's a strong signal you'd be ignoring.

The right option — Reciprocal Rank Fusion (RRF): it combines rankings while ignoring the absolute scores. All that matters is each document's position in each ranking.

The formula

RRF_score(doc) = Σ_i  1 / (k + rank_i(doc))

where:
  k = a constant (typically 60)
  rank_i(doc) = the doc's position in ranking i (1, 2, 3, ...)

If the doc doesn't appear in a ranking, it contributes 0.

The intuition:

• A doc at position #1 contributes 1/(60+1) = 0.0164
• A doc at position #5 contributes 1/(60+5) = 0.0154
• A doc at position #20 contributes 1/(60+20) = 0.0125

The contributions are non-linear — the gap between top-5 and top-20 matters, but the gap between #1 and #2 is tiny. The function rewards documents that consistently rank high across multiple rankings.

The implementation

# rrf.py
from typing import List
from collections import defaultdict


def reciprocal_rank_fusion(
    rankings: List[List[str]],
    k: int = 60,
) -> List[tuple[str, float]]:
    """
    Combine multiple rankings using Reciprocal Rank Fusion.

    Args:
        rankings: List of rankings, each ranking is a list of doc_ids ordered by relevance.
        k: RRF constant (default 60, value used in original paper).

    Returns:
        List of (doc_id, rrf_score) tuples, sorted by score descending.
    """
    rrf_scores = defaultdict(float)

    for ranking in rankings:
        for rank, doc_id in enumerate(ranking, start=1):
            rrf_scores[doc_id] += 1.0 / (k + rank)

    sorted_results = sorted(rrf_scores.items(), key=lambda x: -x[1])
    return sorted_results


# Example usage
ranking_1 = ["doc_a", "doc_b", "doc_c", "doc_d", "doc_e"]      # query 1
ranking_2 = ["doc_b", "doc_a", "doc_f", "doc_c", "doc_g"]      # query 2
ranking_3 = ["doc_c", "doc_a", "doc_b", "doc_h", "doc_d"]      # query 3

fused = reciprocal_rank_fusion([ranking_1, ranking_2, ranking_3])
print("Document : RRF score")
for doc_id, score in fused[:5]:
    print(f"  {doc_id}: {score:.4f}")

Output:

Document : RRF score
  doc_a: 0.0492    ← appears at position 1, 2, 2 (top 3 in all of them)
  doc_b: 0.0476    ← appears at position 2, 1, 3 (top 3 in all of them)
  doc_c: 0.0467    ← appears at position 3, 4, 1
  doc_d: 0.0312    ← appears at position 4, -, 5
  doc_e: 0.0164    ← appears in only one

doc_a wins because it's consistently in the top-3 of all 3 queries. doc_e comes last because it only appears in one.

The complete end-to-end pipeline

# expansion_pipeline.py
import chromadb
from chromadb.utils import embedding_functions
import os

openai_ef = embedding_functions.OpenAIEmbeddingFunction(
    api_key=os.getenv("OPENAI_API_KEY"),
    model_name="text-embedding-3-small",
)
chroma_client = chromadb.PersistentClient(path="./chroma_db")
collection = chroma_client.get_collection("docs", embedding_function=openai_ef)


def search_with_expansion(query: str, top_k: int = 5, num_expansions: int = 5):
    """
    The complete pipeline: expand query → search each → fuse with RRF → final top-K.
    """
    # 1. Expand
    expansion_result = expand_query(query, num_expansions=num_expansions)
    queries_to_search = [query] + expansion_result.queries  # include the original too

    # 2. Search with each query (in parallel if possible)
    all_rankings = []
    for q in queries_to_search:
        results = collection.query(query_texts=[q], n_results=10)
        all_rankings.append(results['ids'][0])

    # 3. Fuse with RRF
    fused = reciprocal_rank_fusion(all_rankings)
    top_ids = [doc_id for doc_id, _ in fused[:top_k]]

    # 4. Fetch the final documents
    final_docs = collection.get(ids=top_ids)
    return final_docs


# Try it
query = "fastapi auth"
results = search_with_expansion(query, top_k=5)

print(f"Query: {query}\n")
print(f"Top 5 after expansion + RRF:")
for i, (doc, doc_id) in enumerate(zip(results['documents'], results['ids']), 1):
    print(f"\n#{i} [{doc_id}]")
    print(f"   {doc[:120]}...")

The optimization: queries in parallel

If you have 5 queries to search, running them sequentially adds 5x the latency. Parallelize:

from concurrent.futures import ThreadPoolExecutor


def parallel_retrieve(queries: list[str], top_k_per_query: int = 10):
    """Runs multiple retrievals in parallel."""
    def search_one(q):
        return collection.query(query_texts=[q], n_results=top_k_per_query)

    with ThreadPoolExecutor(max_workers=5) as executor:
        results_list = list(executor.map(search_one, queries))

    return [r['ids'][0] for r in results_list]

With 5 parallel workers, the 5 retrieval queries run in the time of 1.

When to expand and when not to

Query expansion isn't free: it adds ~600-800ms (1 LLM call + 5 retrieval calls + RRF) and ~$0.001 per query (the LLM call). It isn't always worth it.

When you SHOULD expand

When you should NOT expand

The dynamic skip

Ideally you decide per query whether expanding is worth it:

def should_expand(query: str) -> bool:
    """A simple heuristic for deciding whether to expand."""
    word_count = len(query.split())

    # Very short queries are ambiguous → expand
    if word_count <= 3:
        return True

    # Queries with exact identifiers → don't expand
    has_identifier = any(c in query for c in ['_', '`', '(', ')', '/']) or \
                     any(word.isupper() for word in query.split() if len(word) > 2)
    if has_identifier:
        return False

    # Long, well-formed queries → don't expand
    if word_count > 12 and "?" in query:
        return False

    # Default: expand
    return True


def smart_search(query: str, top_k: int = 5):
    if should_expand(query):
        return search_with_expansion(query, top_k=top_k)
    else:
        results = collection.query(query_texts=[query], n_results=top_k)
        return results

The benefit: ~30-50% of queries don't need expansion. This saves latency and cost on those without losing any quality.

Traps and common mistakes

Trap 1: a prompt without examples leads to weak expansions

The mistake:

Generate 5 search queries related to: "fastapi auth"

The symptom: the LLM generates trivial variations ("fastapi authentication", "fastapi auth tutorial") that add no diversity.

How to prevent it: a prompt with concrete examples (few-shot). LLMs learn far better from examples than from abstract instructions.

Trap 2: temperature=0 generates identical expansions

The mistake: temperature=0 for consistency.

The symptom: the 5 "expansions" are nearly identical. RRF gets no diversity to work with.

How to prevent it: temperature=0.4-0.6 for diversity without drifting from the intent.

Trap 3: ignoring the original query

The mistake:

queries_to_search = expanded_queries  # only the 5 generated ones

The symptom: if the original query was already good, you lose its ranking in the fusion.

How to prevent it: always include the original query alongside the expansions. The original query is just one more "expansion", with equal weight:

queries_to_search = [query] + expanded_queries

Trap 4: too many expansions

The mistake: generating 10-20 expansions "just to be safe".

The symptom:

• Latency explodes (10 retrievals + a longer LLM call).
• Cost rises linearly.
• Quality improves only marginally — 5 expansions already cover the likely interpretations; expansions 10-20 just add noise.

How to prevent it: stay between 4 and 6 expansions. That's the empirical sweet spot.

Trap 5: expansions that change the intent

The mistake: a prompt with no guardrails. The LLM expands "why is FastAPI slow" into queries like "how to make FastAPI faster".

The symptom: the expansion now searches for "how to optimize" when the user wanted to understand "why it's sometimes slow". The system answers the opposite question.

How to prevent it: an explicit instruction in the prompt: "Preserve the user's INTENT — don't add unrelated topics or change the question's stance."

Trap 6: caching without accounting for trivial variants

The mistake: caching the expansion for "fastapi auth" and "FastAPI Auth" as two different queries.

The symptom: a low cache hit rate (~10%) because every capitalization/whitespace variant generates a new expansion.

How to prevent it: normalize the query before caching:

import hashlib

def cache_key(query: str) -> str:
    normalized = query.lower().strip()
    normalized = " ".join(normalized.split())  # collapse whitespace
    return hashlib.md5(normalized.encode()).hexdigest()

The typical hit rate with normalization: 40-60%.

Applied exercise

The scenario: you're an AI Engineer at a company providing technical support for DevOps tools. The log data:

• 10,000 real queries from the last month
• 51% are short/ambiguous queries (1-4 tokens) — "k8s deploy", "docker timeout", "jenkins fail"
• 30% are queries with exact identifiers — "kubectl get pods not working", "ERR_NETWORK_TIMEOUT_504"
• 19% are well-formed queries

The current system: cosine + cross-encoder rerank. The metrics:

• Precision@5: 84%
• Recall@5: 68% (this is the problem — the system isn't finding many relevant docs)

Your job:

1. Decide whether query expansion would help here. Justify it with the data.
2. Design the implementation, including the dynamic skip.
3. Estimate the expected impact and the extra monthly cost (5K queries/day).

def smart_pipeline(query: str, top_k: int = 5):
    """The pipeline with conditional expansion."""

    # The heuristic for deciding whether to expand
    word_count = len(query.split())
    has_identifier = any(c in query for c in ['_', '`', '/']) or \
                     any(re.match(r'[A-Z][A-Z_]+\d*', w) for w in query.split())  # e.g. ERR_TIMEOUT_504
    is_well_formed = word_count > 8 and "?" in query

    if has_identifier or is_well_formed:
        # Skip expansion: direct query + retrieval + rerank
        return search_with_rerank_only(query, top_k=top_k)
    elif word_count <= 5:
        # A short/ambiguous query: use expansion
        return search_with_expansion_and_rerank(query, top_k=top_k)
    else:
        # Default: the direct pipeline
        return search_with_rerank_only(query, top_k=top_k)


def search_with_expansion_and_rerank(query: str, top_k: int = 5):
    # 1. Expand
    expansion = expand_query(query, num_expansions=5)
    queries_to_search = [query] + expansion.queries

    # 2. Parallel retrieval
    all_rankings = parallel_retrieve(queries_to_search, top_k_per_query=15)

    # 3. RRF to combine them
    fused = reciprocal_rank_fusion(all_rankings)
    candidate_ids = [doc_id for doc_id, _ in fused[:30]]

    # 4. Fetch the complete documents
    candidates = collection.get(ids=candidate_ids)

    # 5. Cross-encoder rerank over the 30 candidates
    reranked = cross_encoder_rerank(query, candidates['documents'], top_k=top_k)
    return reranked

queries_per_day = 5000
expansion_rate = 0.51  # 51% get expanded
expanded_queries_per_day = queries_per_day * expansion_rate

# Each expansion: 1 LLM call (~500 tokens in, ~200 tokens out with gpt-4o-mini)
cost_per_expansion = (500 / 1_000_000 * 0.15) + (200 / 1_000_000 * 0.60)
# = $0.000195

monthly_cost = expanded_queries_per_day * 30 * cost_per_expansion
print(f"Monthly expansion cost: ${monthly_cost:.2f}")

Summary and next step

What you learned:

• Query expansion generates several versions of an ambiguous query and combines the results to improve recall.
• The typical gain: +20-30% recall, especially on short or conceptual queries.
• The cost: ~600-800ms of latency + ~$0.001 per query (the LLM call for the expansion).
• Generating quality expansions requires a prompt with examples (few-shot), structured outputs, and temperature=0.4-0.6.
• Reciprocal Rank Fusion combines rankings while ignoring absolute scores. It's the correct technique, unlike naively summing the scores.
• The dynamic skip: don't expand queries with exact identifiers or well-formed ones. It saves ~50% of the extra cost.
• Include the original query in the search set, not just the expansions.
• The common trap: expansions that change the user's intent. An explicit prompt prevents it.

Checkpoint: before moving on, you should be able to:

• Implement query expansion with structured outputs and RRF in under 100 lines.
• Design a dynamic skip with heuristics that detect queries that do NOT need expansion.
• Calculate the extra monthly cost of query expansion for a given volume.

Next capsule: 04 — Query Rewriting.

Query expansion attacks ambiguous queries. Query rewriting attacks incomplete or badly phrased ones. The idea: instead of generating several parallel queries, you transform the original query into a better version (with more context, better structured). It's the complementary technique — and together they cover all three failure modes you saw in M03/02.

Resources

1. Reciprocal Rank Fusion Paper (Cormack et al., 2009) — The original paper
2. LangChain — Multi-Query Retriever — The LangChain implementation
3. LlamaIndex — Query Transformation — Alternative patterns
4. Pinecone — Query Optimization — A practical tutorial
5. OpenAI — Structured Outputs — For a robust implementation
6. Anthropic — Contextual Retrieval — A complementary technique

Estimated time: 30-35 minutes Next: 04-query-rewriting.md