Capsule 04: Query rewriting — transforming the query before you search

Capsule overview

Query expansion (capsule 03) attacks the ambiguity problem: a query with several interpretations becomes several queries, one per interpretation. Query rewriting attacks a different problem: queries that are incomplete, badly phrased or missing context. Instead of generating several parallel versions, you transform the original query into a better version and search with that single improved version.

When you want each technique:

Ambiguous ("fastapi auth")            → expansion (5 parallel queries)
Incomplete ("how to deploy")          → rewriting with the chat's context
Badly phrased ("python async perf")   → structural rewriting into natural language

This capsule teaches you to implement query rewriting in three variants (clarification, contextual, structural), pick the right one for each case, and combine it with other techniques (rewriting + expansion in a cascade, for queries that have both problems).

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

• ✅ Implement query rewriting with an LLM and structured outputs
• ✅ Tell apart three patterns: clarification, rewriting with chat context, structural rewriting
• ✅ Combine rewriting with expansion for queries with multiple problems
• ✅ Decide when rewriting is enough vs when you also need expansion
• ✅ Anticipate the subtlest trap: rewriting that changes the user's intent
• ✅ Implement guardrails that catch problematic rewrites before you use them

Estimated time: 30-35 minutes

The insight: a query optimized for your system, not for Google

As we saw in M03/02, users type queries optimized for Google (short, keyword-style). Your RAG system needs queries optimized for semantic search (complete sentences in natural language). Rewriting is the translation layer.

The user's query: "fastapi auth"
                            │
                            │  Rewriting (LLM)
                            ▼
The rewritten query: "How do I implement authentication in FastAPI applications,
                      including OAuth2, JWT, and session-based methods?"
                            │
                            │  Search (cosine similarity)
                            ▼
Top-K results

The key difference from expansion:

• Expansion: 1 query → 5 queries → 5 searches → fusion
• Rewriting: 1 query → 1 improved query → 1 search

Rewriting is faster (1 retrieval vs 5), cheaper (1 simpler LLM call), but it only works for problems solvable with a single transformation. For genuinely ambiguous queries, expansion wins.

The three rewriting variants

Variant 1: structural clarification

It turns keyword-style queries into complete natural sentences. The simplest form.

# query_rewriting.py
from openai import OpenAI
from pydantic import BaseModel, Field
import os


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


class RewrittenQuery(BaseModel):
    rewritten: str = Field(description="The improved query in natural language")
    reasoning: str = Field(description="One-sentence explanation of what changed")


STRUCTURAL_PROMPT = """You are an expert at rewriting search queries for retrieval systems.

Given a keyword-style or fragmentary query, rewrite it as a complete, natural-language
question. Preserve ALL keywords and the user's intent. Do NOT add information that
wasn't implied in the original.

Examples:

Original: "python async performance"
Rewritten: "How does Python's async/await pattern affect performance in concurrent applications?"

Original: "kubernetes pod restart"
Rewritten: "How do I restart a Kubernetes pod, and what are the implications?"

Original: "fastapi sql injection"
Rewritten: "How can I prevent SQL injection in FastAPI applications?"

Now rewrite the user's query."""


def rewrite_structural(query: str) -> RewrittenQuery:
    """Convert keyword-style query to natural language."""
    response = client.beta.chat.completions.parse(
        model="gpt-4o-mini",
        messages=[
            {"role": "system", "content": STRUCTURAL_PROMPT},
            {"role": "user", "content": f'Original: "{query}"\nRewritten:'},
        ],
        response_format=RewrittenQuery,
        temperature=0.1,  # low for consistency, not creativity
    )
    return response.choices[0].message.parsed


# Try it
test_queries = ["python async performance", "kubernetes pod restart", "fastapi sql injection"]
for q in test_queries:
    result = rewrite_structural(q)
    print(f"Original:  {q}")
    print(f"Rewritten: {result.rewritten}")
    print(f"Reasoning: {result.reasoning}\n")

Variant 2: rewriting with the chat's context

When the user is mid-conversation, their "incomplete" query usually has implicit context in the earlier messages. Use it.

CONTEXTUAL_PROMPT = """You are an expert at rewriting search queries using conversation context.

The user is in a chat about a specific topic. Their previous messages may contain
context that the current query implicitly references.

Given the chat history and the current (potentially incomplete) query, rewrite the
query to be self-contained — including any context that's missing.

Rules:
1. ONLY add context that's clearly implied by the conversation
2. Do NOT add unrelated information
3. Preserve the user's exact intent and tone
4. If the query is already complete, return it unchanged

Examples:

Chat history:
- User: "I'm building a FastAPI app"
- Bot: "Great! What would you like to know?"
Current query: "how to add authentication"
Rewritten: "How do I add authentication to my FastAPI app?"

Chat history:
- User: "I'm using PostgreSQL with SQLAlchemy"
- Bot: "Sounds good. Any specific issue?"
Current query: "the connection keeps dropping"
Rewritten: "Why does my SQLAlchemy connection to PostgreSQL keep dropping?"

Chat history:
- User: "I'm deploying to AWS"
- Bot: "OK, are you using ECS or EKS?"
- User: "ECS"
- Bot: "Got it"
Current query: "how to set environment variables"
Rewritten: "How do I set environment variables in an ECS deployment on AWS?"
"""


def rewrite_with_context(query: str, chat_history: list[str]) -> RewrittenQuery:
    """Rewrite incomplete query using chat history."""
    history_str = "\n".join(chat_history[-5:])  # the last 5 messages

    response = client.beta.chat.completions.parse(
        model="gpt-4o-mini",
        messages=[
            {"role": "system", "content": CONTEXTUAL_PROMPT},
            {
                "role": "user",
                "content": f"Chat history:\n{history_str}\n\nCurrent query: \"{query}\"\nRewritten:",
            },
        ],
        response_format=RewrittenQuery,
        temperature=0.2,
    )
    return response.choices[0].message.parsed


# Try it
chat = [
    "User: I'm using FastAPI with PostgreSQL",
    "Bot: Great, what would you like to know?",
    "User: my queries are slow",
]
result = rewrite_with_context("how to optimize", chat)
print(f"Rewritten: {result.rewritten}")
# Output: "How do I optimize slow PostgreSQL queries in my FastAPI application?"

Variant 3: clarifying rewriting (typos, ambiguity)

For queries with typos, abbreviations, or minor ambiguity that can be resolved with a single dominant interpretation.

CLARIFY_PROMPT = """Rewrite this search query to fix typos, expand abbreviations, and clarify
without changing meaning. If the original is already clear, return it unchanged.

Examples:

Original: "fastpi autentication"   →  "FastAPI authentication"
Original: "k8s pod stuck"           →  "Kubernetes pod stuck"
Original: "psql conn timout"        →  "PostgreSQL connection timeout"
Original: "How to use OAuth2 in FastAPI?" →  "How to use OAuth2 in FastAPI?" (unchanged)

Original query: "{query}"
Rewritten:"""


def rewrite_clarify(query: str) -> str:
    """Fix typos and expand abbreviations."""
    response = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[
            {"role": "user", "content": CLARIFY_PROMPT.format(query=query)},
        ],
        temperature=0.0,  # deterministic, for consistency
        max_tokens=100,
    )
    return response.choices[0].message.content.strip()

Combining rewriting with expansion in a cascade

For queries with MULTIPLE problems (incomplete + ambiguous), combine both techniques:

def smart_query_optimization(query: str, chat_history: list[str] = None) -> list[str]:
    """
    The complete query optimization pipeline:
    1. Rewriting (incomplete → complete, badly phrased → well phrased)
    2. Expansion (if the rewritten query is still ambiguous → several interpretations)
    """
    # Step 1: rewriting, based on the context
    if chat_history and len(chat_history) > 0:
        rewritten = rewrite_with_context(query, chat_history).rewritten
    elif is_keyword_style(query):
        rewritten = rewrite_structural(query).rewritten
    else:
        rewritten = query  # it's already fine

    # Step 2: is the rewritten query still ambiguous?
    if is_ambiguous(rewritten):
        # Apply expansion over the rewritten query
        expansion = expand_query(rewritten, num_expansions=4)
        return [rewritten] + expansion.queries
    else:
        # A single query is enough
        return [rewritten]


def is_keyword_style(query: str) -> bool:
    """A simple heuristic for detecting keyword style."""
    word_count = len(query.split())
    if word_count <= 5 and "?" not in query:
        # Few words + no question mark = probably a keyword query
        return True
    return False


def is_ambiguous(query: str) -> bool:
    """Detect ambiguity after the rewrite."""
    word_count = len(query.split())
    # If it's still short after the rewrite, it's probably ambiguous
    if word_count <= 6:
        return True
    return False

Examples:

Original query: "fastapi auth"        (short + ambiguous)
  Rewriting →  "How do I implement authentication in FastAPI?"  (structured)
  Expansion →  ["How do I implement OAuth2 auth in FastAPI?",
                 "How do I implement JWT auth in FastAPI?", ...]

Original query: "how to optimize"     (incomplete, in a PostgreSQL chat)
  Rewriting with context → "How do I optimize slow PostgreSQL queries?"
  Expansion isn't needed (the query is complete and specific)

Original query: "explain HNSW"        (clear and specific)
  Rewriting → "Explain the HNSW algorithm" (no significant change)
  Expansion isn't needed

When to use each technique

Traps and common mistakes

Trap 1: rewriting that changes the intent

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

The symptom: the system answers "to make FastAPI faster, do X" when the user wanted to understand why it's sometimes slow.

How to prevent it: an explicit instruction in the prompt: "Preserve the user's exact intent. Do NOT change a 'why' question to a 'how to' question."

Trap 2: rewriting with a slow LLM by default

The mistake: using GPT-4 for rewriting when GPT-4o-mini is enough.

The symptom: rewriting adds 800ms when it could add 200ms.

How to prevent it: GPT-4o-mini is plenty for rewriting (a simple structural task). Use GPT-4 only if the measurable quality improves significantly — which is rare.

Trap 3: a chat history that's too long

The mistake: you pass the last 50 chat messages to the rewriter.

The symptom: the context becomes noisy, and the LLM can mix topics from old messages.

How to prevent it: the last 3-5 messages are enough. If the query references something from 20 messages ago, that isn't a rewriting problem — the user should have been more explicit.

Trap 4: applying rewriting to EVERY query

The mistake: you turn rewriting on for all traffic, indiscriminately.

The symptom: queries that were already well formed go through an LLM call for nothing. Significant extra cost.

How to prevent it: detect well-formed queries and skip them. A simple heuristic:

def needs_rewriting(query: str) -> bool:
    word_count = len(query.split())
    has_question_mark = "?" in query
    has_full_sentence_structure = any(query.lower().startswith(w) for w in [
        "how", "what", "why", "when", "where", "which", "can", "should", "is", "do"
    ])

    if word_count > 8 and (has_question_mark or has_full_sentence_structure):
        return False  # already well formed
    return True

Trap 5: structured outputs on a small model version

The mistake: you use gpt-3.5-turbo with structured outputs.

The symptom: the model doesn't support structured outputs and returns malformed JSON.

How to prevent it: structured outputs require gpt-4o-mini or better. For older models, use a prompt + manual parsing with strict validation.

Trap 6: aggressive rewriting of multilingual queries

The mistake: your rewriting prompt is in English. It receives a query in Spanish. The LLM translates it to English "because that's what feels natural given the prompt".

The symptom: the translated query no longer matches the Spanish docs. Recall collapses.

How to prevent it: an explicit instruction, "Preserve the original language of the query":

PROMPT_MULTILINGUAL = """Rewrite the query to natural language.
IMPORTANT: Preserve the original language of the query. If the query is in Spanish, the rewrite must also be in Spanish.
"""

Applied exercise

The scenario: you're an AI Engineer at a customer support company. The log data:

• 8000 queries/day
• 65% are conversational chat queries (with history available)
• 25% are direct queries (with no prior context)
• 10% are queries with typos or technical abbreviations

An analysis of a 200-query log sample shows:

"how deploy app"               →  Incomplete (in a FastAPI chat)
"k8s pod not working"          →  Abbreviation (Kubernetes)
"timeout on my api"            →  Incomplete (in a chat about a specific endpoint)
"how to configure"             →  Incomplete (no context in the chat)
"fastapi async vs sync"        →  Well formed
"why does my code fail"        →  Vague (needs more info)

Your job:

1. Design the rewriting/expansion pipeline given the data.
2. Define the fallback order for when a technique doesn't apply.
3. Estimate the extra monthly cost.

def smart_query_pipeline(query: str, chat_history: list[str] = None) -> str | list[str]:
    """
    A pipeline that applies the right technique for the query's state.
    Returns 1 query (if rewriting is enough) or a list of queries (if expansion is needed too).
    """

    # Step 1: clarify the typos/abbreviations first
    if has_typo_or_abbreviation(query):
        query = rewrite_clarify(query)
        # e.g. "k8s pod not working" → "Kubernetes pod not working"

    # Step 2: if it's in a chat, rewrite with the context
    if chat_history and len(chat_history) >= 2:
        rewritten = rewrite_with_context(query, chat_history).rewritten
        # e.g. "how deploy app" + a FastAPI chat → "How do I deploy a FastAPI app?"

        # If the rewritten query is specific, a single search is enough
        if not is_ambiguous(rewritten):
            return rewritten

        # If it's still ambiguous, add expansion
        expansion = expand_query(rewritten, num_expansions=3)
        return [rewritten] + expansion.queries

    # Step 3: no chat history but keyword-style → structural rewriting
    if is_keyword_style(query):
        rewritten = rewrite_structural(query).rewritten

        if not is_ambiguous(rewritten):
            return rewritten
        else:
            expansion = expand_query(rewritten, num_expansions=4)
            return [rewritten] + expansion.queries

    # Step 4: a well-formed query with no context → is it ambiguous?
    if is_ambiguous(query):
        # Expansion only, no rewriting
        expansion = expand_query(query, num_expansions=4)
        return [query] + expansion.queries

    # Step 5: a well-formed, specific query → straight through
    return query

1. Does it have typos? → rewrite_clarify (fast, almost always applicable)
2. Is it in a chat with history? → rewrite_with_context (high priority if the chat exists)
3. Is it keyword-style? → rewrite_structural
4. Is it still ambiguous after rewriting? → add expansion
5. Otherwise → the direct query

QUERIES_PER_DAY = 8000
DAYS_PER_MONTH = 30

# The distribution by processing type
PCT_NEEDS_REWRITING = 0.50    # ~50% needs rewriting
PCT_NEEDS_EXPANSION = 0.20    # ~20% also needs expansion
PCT_NO_OPTIMIZATION = 0.30    # ~30% go straight through

# The cost per LLM call (gpt-4o-mini)
COST_PER_CALL_REWRITING = 0.0001  # rewriting is a short prompt
COST_PER_CALL_EXPANSION = 0.0002  # expansion is a longer prompt

# The cost per extra retrieval (assumes 4 extra retrievals on the queries with expansion)
COST_PER_RETRIEVAL_EXTRA = 0.0001  # the OpenAI embedding of the query

monthly_cost = (
    QUERIES_PER_DAY * DAYS_PER_MONTH * PCT_NEEDS_REWRITING * COST_PER_CALL_REWRITING +
    QUERIES_PER_DAY * DAYS_PER_MONTH * PCT_NEEDS_EXPANSION * COST_PER_CALL_EXPANSION +
    QUERIES_PER_DAY * DAYS_PER_MONTH * PCT_NEEDS_EXPANSION * 4 * COST_PER_RETRIEVAL_EXTRA
)

print(f"Extra monthly cost: ${monthly_cost:.2f}")

Summary and next step

What you learned:

• Query rewriting transforms one query into an improved version, unlike query expansion, which generates several parallel versions.
• Three variants: structural (keyword → natural language), contextual (with chat history), clarifying (typos/abbreviations).
• Combine rewriting + expansion in a cascade when the query has multiple problems.
• The typical gain: +10-15% precision with rewriting, +20-30% recall with expansion. Combined, you get both.
• The main trap: rewriting that changes the intent. An explicit prompt + manual sampling prevent it.
• The dynamic skip: well-formed queries don't need rewriting. Detecting and skipping them saves cost and latency.
• For multilingual: an explicit instruction to preserve the original language.

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

• Implement all three rewriting variants with structured outputs.
• Design a cascading pipeline that combines rewriting + expansion as needed.
• Detect queries that do NOT need rewriting with simple heuristics.

Next capsule: 05 — Query Decomposition.

Rewriting and expansion attack short or ambiguous queries. Capsule 05 covers the opposite case: complex queries that pack several questions into one. E.g. "how do I deploy FastAPI with Docker, configure SSL, and monitor with Prometheus?" — those are three distinct questions requiring three different sets of chunks. Decomposition splits complex queries into sub-queries and combines the results.

Resources

1. LangChain — Query Construction Guide — Advanced rewriting patterns
2. Anthropic — Multi-Step Question Answering — When rewriting alone isn't enough
3. Microsoft — Query Rewriting in Production — A real case with Azure Cognitive Search
4. OpenAI — Structured Outputs Guide — For a robust implementation
5. Pinecone — Query Optimization Series — The complete tutorial
6. Stanford NLP — Query Reformulation Theory — Foundational

Estimated time: 30-35 minutes Next: 05-query-decomposition.md