RAG Query Expansion

RAG query expansion turns one user question into several retrieval-oriented queries, searches with each, and combines the results before generation. It helps when the user's wording does not match the terminology in your documents, and you can implement it with a small language-model call plus deterministic rank fusion. The safest production design keeps the original query, limits expansion count, deduplicates results, and measures retrieval quality separately from answer quality.

Why rag query expansion improves retrieval

Retrieval-augmented generation fails surprisingly often before the generator sees any useful context. A user asks, "Why does checkout fail after five minutes?" while the incident report says "payment session TTL expiration." A single embedding may connect those phrases, but it may rank a generic checkout guide above the exact incident. Keyword search can miss the document entirely.

Expansion generates alternate search expressions that expose hidden vocabulary. For the example above, useful variants could include payment session expiration checkout, checkout token TTL, and cart payment session timeout. Each query creates another route to relevant chunks.

This technique addresses four common retrieval problems:

• Vocabulary mismatch, users and authors name the same concept differently.
• Underspecified questions, important entities or constraints are implicit.
• Multi-aspect questions, one query embedding blends distinct requirements.
• Acronyms and aliases, such as OIDC, OpenID Connect, and identity federation.

Expansion is not a license to invent facts. The expansion model proposes search strings, not answers. Retrieved documents remain the evidence, and the original question remains the generator's task.

Build a minimal rag query expansion pipeline

The following example uses the OpenAI Python SDK for query generation and a local BM25 index for retrieval. It is intentionally small enough to run as one file. Use environment variables for credentials and pin dependencies in your real application.

Install the packages:

python -m venv .venv
source .venv/bin/activate
python -m pip install openai rank-bm25
export OPENAI_API_KEY="your-key"

Save this as expand_search.py:

import json
import os
import re
from dataclasses import dataclass

from openai import OpenAI
from rank_bm25 import BM25Okapi


@dataclass(frozen=True)
class Document:
    id: str
    text: str


DOCS = [
    Document("runbook-17", "Payment sessions have a five minute TTL. "
             "Expired sessions return CHECKOUT_SESSION_EXPIRED."),
    Document("guide-4", "The checkout page retries transient gateway errors twice."),
    Document("api-9", "Create a new payment session when the previous token expires."),
    Document("ops-2", "Cart data is retained for thirty days after creation."),
]


def tokenize(text: str) -> list[str]:
    return re.findall(r"[a-z0-9_]+", text.lower())


def expand_query(client: OpenAI, question: str) -> list[str]:
    response = client.responses.create(
        model=os.environ.get("EXPANSION_MODEL", "gpt-4.1-mini"),
        input=[
            {
                "role": "system",
                "content": (
                    "Generate exactly three compact search queries for a technical "
                    "document index. Preserve product names, error codes, versions, "
                    "and numbers. Cover distinct plausible terminology. Return JSON "
                    "with one key named queries and no explanation."
                ),
            },
            {"role": "user", "content": question},
        ],
        text={"format": {"type": "json_object"}},
    )
    data = json.loads(response.output_text)
    candidates = [question, *data.get("queries", [])]
    return list(dict.fromkeys(q.strip() for q in candidates if q.strip()))[:4]


def search(index: BM25Okapi, query: str, limit: int = 3) -> list[tuple[str, float]]:
    scores = index.get_scores(tokenize(query))
    order = sorted(range(len(DOCS)), key=lambda i: scores[i], reverse=True)
    return [(DOCS[i].id, float(scores[i])) for i in order[:limit]]


def rrf(result_lists: list[list[tuple[str, float]]], k: int = 60) -> list[tuple[str, float]]:
    fused: dict[str, float] = {}
    for results in result_lists:
        for rank, (doc_id, _) in enumerate(results, start=1):
            fused[doc_id] = fused.get(doc_id, 0.0) + 1.0 / (k + rank)
    return sorted(fused.items(), key=lambda item: item[1], reverse=True)


def main() -> None:
    client = OpenAI()
    index = BM25Okapi([tokenize(doc.text) for doc in DOCS])
    question = "Why does checkout fail after five minutes?"
    queries = expand_query(client, question)
    fused = rrf([search(index, query) for query in queries])
    print(json.dumps({"queries": queries, "results": fused}, indent=2))


if __name__ == "__main__":
    main()

Run it:

python expand_search.py

The original question is deliberately inserted first. If generation returns malformed, redundant, or unhelpful variants, retrieval still has the user's exact words. dict.fromkeys removes exact duplicates while preserving order. In production, also normalize whitespace and compare lowercase forms.

Reciprocal rank fusion, or RRF, combines rankings without assuming that scores from different queries are comparable. BM25 scores vary with query terms, and vector similarities vary across indexes and configurations. RRF only uses rank positions, so it is robust and easy to inspect.

Choose an expansion strategy

No single prompt fits every corpus. Choose the smallest strategy that targets observed retrieval failures.

Paraphrase expansion rewrites the entire question with synonyms. It works well for natural-language knowledge bases and support articles. Ask for distinct terminology, because superficial word-order changes add cost without improving recall.

Acronym and entity expansion preserves the original entity while adding aliases. For S3 403 after enabling SSE-KMS, expansions might include Amazon S3 AccessDenied KMS key policy and SSE-KMS object access IAM. This is useful for technical corpora with product-specific language.

Subquery decomposition splits a compound request into independently searchable parts. For "Compare blue-green and canary deployments, including rollback and database migration risks," retrieve deployment definitions, rollback mechanics, and database compatibility separately. Decomposition is more powerful than paraphrasing, but it can over-retrieve and dilute context.

HyDE-style expansion asks a model to draft a hypothetical relevant passage, embeds that passage, and retrieves similar real passages. It can help dense retrieval because document-like prose resembles indexed chunks. Never present the hypothetical passage as evidence, and do not use its unsupported details as filters.

Metadata-aware expansion emits a text query plus explicit constraints such as service, language, version, or date. This works only when metadata is reliable. Prefer validated structured output and allowlisted fields rather than parsing free-form text.

A practical default is two or three paraphrases plus the original query. Add decomposition only for questions containing multiple explicit tasks. Use HyDE after evaluation shows that short-query embeddings are a recurring weakness.

Use structured output and validation

Treat model output as untrusted input. A good schema separates queries from filters and gives every expansion a purpose. The exact SDK helper for schema validation can change, so the portable approach is to request JSON and validate it locally.

from dataclasses import dataclass
from typing import Any


ALLOWED_FILTERS = {"service", "version", "language"}


@dataclass(frozen=True)
class ExpansionPlan:
    queries: list[str]
    filters: dict[str, str]


def validate_plan(value: Any, original: str) -> ExpansionPlan:
    if not isinstance(value, dict):
        raise ValueError("plan must be an object")

    raw_queries = value.get("queries", [])
    if not isinstance(raw_queries, list):
        raise ValueError("queries must be a list")

    queries = [original]
    for item in raw_queries:
        if not isinstance(item, str):
            continue
        cleaned = " ".join(item.split())
        if 2 <= len(cleaned) <= 240 and cleaned.lower() not in {
            q.lower() for q in queries
        }:
            queries.append(cleaned)
        if len(queries) == 4:
            break

    raw_filters = value.get("filters", {})
    filters = {}
    if isinstance(raw_filters, dict):
        for key, item in raw_filters.items():
            if key in ALLOWED_FILTERS and isinstance(item, str) and len(item) <= 80:
                filters[key] = item

    return ExpansionPlan(queries=queries, filters=filters)

Do not let generated filters silently exclude evidence. If a user mentions version 3, a version=3 constraint is defensible. If the model merely guesses version 3, use it as a query term or omit it. Apply authorization filters from trusted application state after expansion, never from the language model.

Prompt injection can also exist in the user's query. The expansion call should not have tools, secrets, or access to retrieved content. Its narrow job is producing search variants. Limit input length, log rejected output, and fall back to the original query on any exception.

Combine lexical and vector retrieval

Query expansion becomes more useful when each variant searches both lexical and vector indexes. Lexical retrieval catches exact error codes, configuration keys, and identifiers. Vector retrieval catches paraphrases and conceptual similarity.

Run the retrieval calls concurrently, then fuse their rankings. This reduces latency compared with sequential searches:

import asyncio
from collections.abc import Awaitable, Callable


Result = list[tuple[str, float]]


async def retrieve_all(
    queries: list[str],
    lexical_search: Callable[[str], Awaitable[Result]],
    vector_search: Callable[[str], Awaitable[Result]],
) -> list[tuple[str, float]]:
    jobs = []
    for query in queries:
        jobs.append(lexical_search(query))
        jobs.append(vector_search(query))

    rankings = await asyncio.gather(*jobs)
    return rrf(rankings)[:20]

Fetch more candidates than the generator will receive. For example, fuse to 20 candidates, rerank those, deduplicate overlapping chunks, then pass a compact final set. The exact counts depend on chunk size, corpus redundancy, context budget, and latency targets, so tune them with evaluation rather than copying fixed values.

Deduplication should operate at more than one level:

• Exact chunk ID removes repeated hits across queries and retrievers.
• Canonical document ID prevents one long document from monopolizing context.
• Text similarity catches duplicated content published under different URLs or versions.
• Metadata rules prefer the requested or latest supported version when the user is ambiguous.

Keep source diversity intentional. A troubleshooting answer may benefit from a runbook, an API reference, and a resolved incident. Ten adjacent chunks from one guide often provide less coverage.

Control latency and cost

Expansion adds one model call and multiplies search operations. Most search backends handle several parallel queries cheaply, but tail latency can still grow. Use a budgeted pipeline:

1. Classify whether expansion is needed using deterministic signals first.
2. Generate at most a small number of distinct variants.
3. Search all variants concurrently with per-request timeouts.
4. Stop or degrade gracefully when the latency budget is exhausted.
5. Cache safe expansion results for normalized, repeated questions.

Skip expansion for exact identifiers, quoted strings, known error codes, and navigational queries when direct search already performs well. Consider expansion for short conceptual questions, low-confidence first-pass retrieval, or questions containing multiple clauses.

A two-stage approach avoids paying the expansion cost on easy traffic. Search the original query first. If the top results pass a calibrated confidence rule, continue directly. Otherwise expand and search again. Do not use an arbitrary similarity threshold across models or indexes. Calibrate it from labeled traffic and monitor it after embedding or corpus changes.

Timeouts need explicit fallbacks. If expansion fails, retrieve with the original query. If one backend times out, fuse available rankings. If reranking fails, use fused order. These fallbacks keep availability independent from an optional quality layer.

Evaluate rag query expansion correctly

Measure retrieval before measuring generated answers. End-to-end answer scores can hide the cause of failure, because a generator might answer from prior knowledge despite poor retrieval or produce a weak answer despite excellent evidence.

Create a query set from real search logs, support tickets, and engineering questions. Remove sensitive data and label the chunks or documents that contain sufficient evidence. Include hard categories such as acronyms, versioned APIs, vague symptoms, compound questions, and exact identifiers.

Useful retrieval metrics include:

• Recall at k, whether any relevant evidence appears in the first k results.
• Mean reciprocal rank, how early the first relevant result appears.
• Precision at k, how much of the candidate set is relevant.
• Context coverage, whether all required aspects of a compound question are present.
• No-answer correctness, whether the system avoids fabricating evidence when the corpus lacks it.

Compare at least four configurations: original query only, original plus paraphrases, hybrid retrieval without expansion, and hybrid retrieval with expansion. Hold chunking, index contents, filters, and reranking constant. Report results by query category because an average can conceal regressions on exact identifiers.

Add assertions to a small offline harness:

def recall_at_k(ranked_ids: list[str], relevant_ids: set[str], k: int) -> float:
    return float(bool(set(ranked_ids[:k]) & relevant_ids))


def evaluate(cases, retrieve, k=5):
    scores = []
    for case in cases:
        ranked = retrieve(case["question"])
        scores.append(recall_at_k(ranked, set(case["relevant_ids"]), k))
    return sum(scores) / len(scores) if scores else 0.0

Do not invent a universal improvement target. Establish a baseline on your data, define acceptable latency and quality gates, and ship only when gains survive category-level review. Re-run the suite after changing the model, prompt, embedding model, tokenizer, chunking, or corpus.

Observe production behavior

Offline labels become stale, so add production signals. Log safely, with redaction and retention controls appropriate for your environment.

Capture the original query, generated variants, prompt version, model identifier, expansion duration, searches attempted, backend timeouts, matched chunk IDs, fusion contributions, reranker order, and final citations. Use stable trace IDs to connect retrieval and generation without putting sensitive text into metric labels.

Monitor distributions rather than a single average:

• Expansion count and duplicate rate.
• Original-query versus expansion-only evidence wins.
• Empty-result and timeout rates.
• Candidate count before and after deduplication.
• Retrieval and end-to-end latency percentiles.
• User reformulation, citation opening, and explicit feedback rates.

Sample traces where expansions dominate the fused ranking. Review whether they reveal useful terminology or semantic drift. Drift patterns should become prompt constraints, validation rules, or evaluation cases.

Roll out with a feature flag. Start with shadow evaluation, where expanded retrieval runs but does not affect answers. Then expose a small traffic slice, compare guardrail metrics, and keep an immediate switch back to original-query retrieval.

Common failure modes

Semantic drift occurs when a variant changes user intent. Preserve named entities, negation, versions, dates, units, and error codes. Rerank against the original question and discard expansions with obvious contradictions.

Query explosion occurs when decomposition creates too many searches and candidates. Cap variants and subqueries, use concurrency limits, and allocate a total retrieval deadline.

Context dilution occurs when higher recall floods the generator with marginal chunks. Rerank, deduplicate, enforce per-document caps, and select evidence that covers distinct aspects.

Filter hallucination occurs when the expansion model invents metadata. Only derive strict filters from explicit user text or trusted application state. Validate names and values against allowlists.

Version mixing occurs when expansions retrieve old and current documentation together. Index version metadata, apply explicit user constraints, and label chunks clearly in the generation context.

Evaluation leakage occurs when prompts are tuned against a tiny fixed test set. Maintain a held-out set, refresh examples from production failures, and review changes by category.

Untraceable relevance occurs when fused results lose the reason they ranked. Store per-query and per-retriever provenance so engineers can reproduce decisions.

Production checklist

• Always include the original query.
• Cap generated variants and their length.
• Preserve identifiers, entities, negation, and version constraints.
• Validate structured output and allowlist metadata fields.
• Apply authorization independently of generated filters.
• Run searches concurrently with deadlines.
• Fuse rankings with a score-independent method such as RRF.
• Deduplicate chunks and limit document concentration.
• Rerank candidates against the original question.
• Keep hypothetical content out of evidence and citations.
• Trace expansion, retrieval, fusion, reranking, and selection.
• Evaluate recall, ranking, coverage, latency, and no-answer behavior.
• Provide fallbacks for every optional stage.
• Roll out behind a feature flag and monitor category-level regressions.

FAQ

How many expanded queries should I generate?

Start with two or three variants plus the original. Increase only when labeled evaluation shows additional recall worth the latency and noise. Distinctness matters more than count.

Should I expand every RAG question?

No. Exact error codes, identifiers, and already-specific navigational queries often work better with direct retrieval. Expand when first-pass confidence is low or when vocabulary mismatch and multiple aspects are likely.

Is query expansion the same as query rewriting?

Rewriting usually replaces a conversational or ambiguous query with one standalone query. Expansion retains the original and adds multiple alternate queries. A production pipeline can first rewrite conversation-dependent context, then expand the standalone result.

Does rag query expansion require a language model?

No. Synonym dictionaries, acronym maps, spelling correction, taxonomy aliases, and relevance feedback can expand queries deterministically. Language models are useful when terminology and intent require contextual interpretation.

Should expanded queries be shown to the answer model?

Usually not. They are retrieval instructions, not evidence. Give the answer model the original question, conversation context if needed, and selected source chunks. Retain expansions in traces for debugging.

What is the best fusion algorithm?

RRF is a strong default because it combines rankings without score calibration. Weighted score fusion can work when scores are normalized and calibrated, but it is easier to misconfigure across retrievers and query variants.

Can expansion make answers worse?

Yes. Drift, obsolete versions, weak filters, and context dilution can push irrelevant evidence upward. Keep the original query, cap variants, rerank against original intent, and require evaluation gains before rollout.

Where should access-control filters run?

They should run as mandatory constraints in the retrieval layer, based on trusted identity and application state. Never let an expansion model loosen, replace, or infer permissions.

When should I use HyDE?

Use it when dense retrieval consistently struggles with very short queries and document-like hypothetical text improves labeled recall. Keep the hypothetical passage separate from retrieved evidence and test for fabricated terminology that causes drift.

How do I know whether expansion helped a specific answer?

Inspect provenance. If necessary evidence was absent for the original query but ranked well for a validated expansion, expansion helped retrieval. Confirm that the final answer cites that evidence and remains faithful to the user's original intent.