Build a Query Rewriting and Expansion Engine for RAG


data-engineering performance databases

AI System Design Deep Dive

Query Rewriting and Expansion Engine

The user's question is not the query your index wants. Something has to translate, and it has 45ms.

⏱ 14 min read 📐 Advanced 🔍 RAG Systems

A reference librarian and a card catalog do very different jobs. The catalog answers the query it was given. The librarian listens to “I need the thing about the tax rule for the other office” and, without asking anything, walks to the shelf that actually holds it. Everything interesting happens in the gap between what the person said and what the index can match. A retrieval system with no rewriting layer is a card catalog with a chat interface bolted on the front.

That gap is measurable. On a production support-assistant corpus of 480 million chunks serving 12,000 queries per second at peak, we measured recall@8 at 0.62 when we embedded the raw user turn and searched a single dense index. The failure distribution was not uniform. Single-entity factual questions retrieved fine, around 0.88. Anaphoric follow-ups, which were 31% of all turns, retrieved at 0.24. Keyword-shaped queries of two or three words retrieved at 0.41, because a three-token query embeds to a vector that sits in a dense, low-information region where three hundred chunks all score between 0.70 and 0.78. Multi-constraint questions that needed two different documents retrieved one of them and stopped.

The naive fixes fail in instructive ways. Making the embedding model bigger moves recall@8 by two or three points and costs 4x per query, because the problem is not embedding quality, it is that the query and the document are written in different registers: users ask questions, documents make statements. Lowering the similarity threshold and retrieving top-50 instead of top-8 raises recall at the cost of precision, and the generator’s answer quality falls off a cliff once the context window holds more distractors than evidence. Asking a large model to rewrite every query buys the recall but adds 400ms to TTFT and $180,000 a year in rewrite tokens at this volume, which is more than the entire retrieval tier costs.

So the engine has to be surgical. It needs to decide cheaply whether a query needs rewriting at all, resolve references against conversation history without a round trip to a frontier model, generate a small set of query forms that probe the index from different angles, run them in parallel against indexes with different failure modes, and merge the results into one ranked list that is better than any single list. We need to solve for recall lift, added latency under 45ms, and cost per query that stays inside the noise of the generation call simultaneously, and any design that optimizes one of those three in isolation is not deployable.

Requirements and Constraints

Functional Requirements

  • Classify every incoming query into a rewrite strategy: passthrough, coreference resolution, decomposition, expansion, or a combination
  • Resolve pronouns, ellipsis, and implicit entities against the last N conversation turns so the rewritten query is self-contained
  • Generate 1 to 4 query variants per request, each targeting a different retrieval mode (lexical, dense, learned sparse)
  • Optionally synthesize a HyDE pseudo-document and embed it as an additional dense probe
  • Expand entities with synonyms, acronyms, and product aliases from a curated lexicon before lexical search
  • Fan out variants across BM25, dense HNSW, and SPLADE indexes in parallel and fuse the rank lists
  • Deduplicate and diversify the fused candidate set, then rerank with a cross-encoder
  • Cache rewrites keyed on a normalized query hash, with a semantic fallback lookup
  • Emit a full trace for every rewrite: variants generated, per-index rank lists, fused scores, and final selection
  • Degrade to the original query on any timeout, validation failure, or quota exhaustion

Non-Functional Requirements

  • Latency: rewrite plus expansion adds no more than 45ms at p99 to the retrieval path. Intent classification under 5ms at p99. Total retrieval path (rewrite plus fan-out plus fusion plus rerank) under 180ms at p99, so TTFT stays under 700ms end to end
  • Throughput: 12,000 queries/second peak, of which roughly 38% are classified as needing a model-generated rewrite, so about 4,560 rewrites/second sustained
  • Fan-out budget: no more than 12 index calls per user query. At 12,000 qps that is 144,000 index calls/second, which sets the shard count
  • Quality: recall@8 at or above 0.85 on the golden set, nDCG@10 at or above 0.72, and no regression greater than 1.5 points on any query class between deploys
  • Cost: rewrite and expansion together stay under $0.40 per 1,000 queries, which is under 8% of the generation cost at our average answer length
  • Capacity: rewrite model fits in 24GB of GPU memory with a batch of 32 concurrent sequences. Rewrite cache holds 40 million entries at roughly 8GB per region
  • Availability: the rewrite tier is not on the critical path for correctness. A total rewrite outage degrades recall, it does not return errors

Constraints

We assume a decoder-only 3B model with a LoRA rewriter adapter for generation, a 1024-dimension embedding model with a float16 HNSW index, and bge-reranker-v2-m3 as the cross-encoder. Hardware is L4 GPUs at 24GB for rewriting and reranking, since neither workload needs A100-class memory bandwidth. Deliberately out of scope: index construction and chunking (that is upstream), the generation call itself, and personalization of retrieval by user history beyond the current session.

High-Level Architecture

The engine is six components, and the ordering between them is the part that carries the design.

Architecture of a query rewriting and expansion engine showing intent gate, anaphora resolver, rewrite cache, rewrite model pool, expansion planner, three retrieval indexes, RRF fusion, and cross-encoder reranking

A request arrives at the RAG API gateway carrying the raw query, a session_id, and a tenant identifier. The intent gate runs a 4ms DistilBERT classifier that emits a rewrite strategy and a confidence. If the strategy is passthrough, which happens for 62% of traffic, the query goes straight to retrieval and the entire rewrite tier is skipped. Otherwise the anaphora resolver pulls the last six turns from the session store and produces a self-contained query. That resolved query is hashed and probed against the rewrite cache; a hit returns the stored variants immediately and bypasses the GPU entirely.

On a miss, the rewrite model pool generates variants under a hard 45ms deadline using constrained JSON decoding on a batched 3B model with a rewriter adapter. The expansion planner takes those variants, attaches lexicon expansions for any recognized entities, optionally requests a HyDE passage, and assigns each variant to the indexes where it will perform best. Variants then fan out in parallel to the BM25 shards, the dense HNSW index, and the SPLADE index.

The returned rank lists converge on the fusion and reranking layer, which applies reciprocal rank fusion across the lists, collapses near-duplicates with maximal marginal relevance, and passes 40 surviving candidates to a cross-encoder that produces the final top-8. Everything the engine did is written to the query telemetry sink, because the traces are the training data for the next version of the rewriter.

Key Insight

The classifier that decides whether to rewrite is more valuable than the rewriter itself. Rewriting a query that was already well formed reduces recall, because a good query is already near its documents in embedding space and any paraphrase moves it away.

The Intent Gate

The intent gate decides what kind of surgery a query needs, and it has to decide in under 5ms without a GPU.

Here is what a smart engineer usually assumes and gets wrong: that rewriting is always safe, so you should just rewrite everything and skip the classifier. We shipped that version first. recall@8 went from 0.62 to 0.79 overall, but it went down on the well-formed query class, from 0.88 to 0.81. The reason is mechanical. A user who types “EMEA data residency requirements for customer PII” has written something that already looks like a sentence out of the target document. Any paraphrase is a random walk away from a point that was already close. Rewriting helps exactly when the query is far from its documents, and hurts when it is near.

So the gate is a five-way classifier over a distilbert-base encoder fine-tuned on 240,000 labeled queries, running on CPU with ONNX Runtime and dynamic int8 quantization. It emits one of passthrough, coref, decompose, expand, or coref_decompose.

# Intent gate: decides the rewrite strategy for a query in under 5ms on CPU.
import numpy as np
import onnxruntime as ort
from tokenizers import Tokenizer

STRATEGIES = ["passthrough", "coref", "decompose", "expand", "coref_decompose"]

class IntentGate:
    def __init__(self, model_path: str, tokenizer_path: str, max_len: int = 64):
        opts = ort.SessionOptions()
        opts.intra_op_num_threads = 2
        opts.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
        self.session = ort.InferenceSession(
            model_path, opts, providers=["CPUExecutionProvider"]
        )
        self.tok = Tokenizer.from_file(tokenizer_path)
        self.tok.enable_truncation(max_length=max_len)
        self.tok.enable_padding(length=max_len)
        self.max_len = max_len

    def classify(self, query: str, turn_index: int) -> tuple[str, float]:
        enc = self.tok.encode(query)
        ids = np.array([enc.ids], dtype=np.int64)
        mask = np.array([enc.attention_mask], dtype=np.int64)
        logits = self.session.run(None, {"input_ids": ids, "attention_mask": mask})[0][0]
        probs = np.exp(logits - logits.max())
        probs /= probs.sum()
        idx = int(probs.argmax())
        conf = float(probs[idx])

        # First turn of a session cannot contain an anaphor with an antecedent.
        if turn_index == 0 and STRATEGIES[idx].startswith("coref"):
            fallback = "decompose" if probs[2] > probs[3] else "expand"
            return fallback, float(max(probs[2], probs[3]))

        # Low confidence defaults to the cheap, safe path.
        if conf < 0.55:
            return "passthrough", conf
        return STRATEGIES[idx], conf

The turn_index == 0 guard is not a detail. Coreference is the single most confident class for the model because anaphors have obvious surface markers, and on the first turn of a session those markers are false positives: “what about the pricing” on turn zero has no antecedent to resolve, and a resolver handed no history will hallucinate one. We measured a 3.4% false coref rate on first turns before adding the guard, and each of those produced a rewrite that invented an entity the user never mentioned.

The confidence floor of 0.55 routes ambiguity to passthrough rather than to a rewrite. That asymmetry is deliberate: a missed rewrite costs recall on one query, while a bad rewrite can retrieve confidently wrong documents that the generator then cites.

Key Insight

The gate's errors are not symmetric, so its threshold should not be either. A false negative loses recall on one query; a false positive rewrites a good query into a worse one and the generator answers confidently from the wrong documents.

Watch Out

Training the gate on query text alone produces a classifier that cannot tell "what about the other one" (needs coref) from "what about the other one, meaning the EU entity" (already self-contained). Concatenate a truncated turn count and the presence of a capitalized entity as features, or accept a 6% error rate on the coref class.

The Anaphora Resolver

The resolver’s job is to make the query self-contained, so that everything downstream can treat it as a standalone search string with no session context attached.

The naive approach is to concatenate the last few turns and embed the whole thing. This is worse than useless. Concatenating six turns of a support conversation produces 600 tokens of mostly irrelevant text, and the resulting embedding is dominated by the conversational filler rather than the current intent. We measured recall@8 on anaphoric queries at 0.24 for raw query, 0.31 for concatenated history, and 0.81 for resolved query. The resolution has to be an actual rewrite, not a context dump.

We run resolution as a separate, smaller task from variant generation, on the same GPU pool but with a different LoRA adapter. The reason to separate them is that resolution is a nearly deterministic, high-precision task with a short output, while variant generation is a creative task with a longer output. Merging them into one prompt makes the model trade off between them, and in our A/B the merged version resolved correctly 91% of the time versus 97% for the dedicated adapter.

# Anaphora resolution: rewrite a follow-up turn into a self-contained query.
from dataclasses import dataclass

RESOLVE_SYSTEM = (
    "Rewrite the final user turn so it is understandable with no conversation "
    "history. Replace pronouns and elliptical references with the entities they "
    "refer to. Change nothing else. Do not answer. Output only the rewritten query."
)

@dataclass
class Turn:
    role: str
    text: str

def build_resolve_prompt(turns: list[Turn], query: str, max_ctx_tokens: int = 320) -> str:
    # Walk backwards so the most recent turns survive truncation.
    kept, budget = [], max_ctx_tokens
    for turn in reversed(turns[-6:]):
        cost = len(turn.text) // 4 + 4
        if cost > budget:
            break
        kept.append(turn)
        budget -= cost
    kept.reverse()
    history = "\n".join(f"{t.role}: {t.text}" for t in kept)
    return (
        f"<|system|>\n{RESOLVE_SYSTEM}\n"
        f"<|history|>\n{history}\n"
        f"<|query|>\n{query}\n"
        f"<|rewritten|>\n"
    )

def is_resolution_safe(original: str, resolved: str, history_text: str) -> bool:
    """Reject resolutions that invent entities absent from the conversation."""
    orig_tokens = set(original.lower().split())
    hist_tokens = set(history_text.lower().split())
    new_tokens = set(resolved.lower().split()) - orig_tokens - hist_tokens
    content = [t for t in new_tokens if len(t) > 3 and t.isalpha()]
    # Allow up to 2 genuinely new function-ish words, reject anything more.
    return len(content) <= 2

Premium Content

Unlock the full article along with everything else in the archive — all in one place.

In-depth analysis Expert insights Full archive access
Unlock Full Article