Retrieval and Query Rewriting: Fix RAG Before It Fails

Your RAG bot has the right document sitting in its index, and it still gives a wrong answer. You check the embeddings, the chunking, the vector database, and everything looks fine. So what went wrong? Very often the problem is the query itself. Today we fix that.

• Index is fine
• Query is the problem

In the last episode we split documents into chunks that are worth retrieving. Now we look at the other side of the match: the question. We will cover the knobs on retrieval, a diversity trick called maximal marginal relevance, and a family of techniques that transform the query before searching. We will also implement the key pieces in plain Python.

• Episode 17
• Retrieval and query rewriting

Here is the plan. First, why user questions make poor search queries. Then the retrieval knobs: top k, thresholds, filters, and diversity. Then the transformations: rewriting, multi query, hypothetical documents, step back, decomposition, and routing. We finish with cost, code, and the mistakes that bite in production.

• Knobs and MMR
• Query transformations
• Code, cost, mistakes

A user question and the document that answers it often look nothing alike. The user asks, how do I get my money back. The document says, refunds are issued within fourteen days of cancellation. The question is short, informal, and phrased as a request. The passage is formal and states facts. Embeddings bridge some of that gap, but not all of it.

• Questions ask
• Documents state

Queries fail in three common ways. They can be vague, like a single word. They can depend on context, like, what about the second one, which means nothing alone. Or they can bundle several questions into one, so the embedding lands between topics and matches none of them well. Each failure needs a different fix.

• Vague or short
• Missing context
• Several questions in one

Picture the embedding space from episode eleven. Documents are dots, and the query is one more dot. Retrieval returns the nearest dots. If the query dot lands in the wrong neighborhood, because it is vague or mixed, the nearest neighbors are wrong, no matter how good the index is. Our job is to move that query dot to the right place.

• Query dot vs document dots
• Nearest is not always right

The first knob is top k, how many chunks you fetch. Suppose the best chunk sits at rank six. With k equal to three, you never see it, and the model has no chance. With k equal to eight, it is in the context. But a larger k adds noise, more tokens to pay for, and more distractors for the model to be confused by.

• Too small: miss the answer
• Too large: noise and cost

A good habit is to separate two numbers: how many candidates you retrieve, and how many you finally put in the prompt. Retrieve generously, say twenty or thirty, then narrow down to the handful that go to the model. Recall is cheap at the retrieval stage. The narrowing step is where precision comes from, and we will improve it a lot in the next episode.

• Retrieve 20 to 50
• Keep 3 to 5 for the prompt

The second knob is a similarity threshold. Instead of always returning the top k, you drop anything below a minimum score. If the scores are zero point eight two, zero point seven nine, then zero point six one, and zero point five five, a threshold of zero point six keeps three. The big benefit: if nothing passes, you know the corpus probably has no answer, and you can say so.

• Drop weak matches
• Enables 'no answer'

A warning: a score of zero point six means different things for different embedding models. There is no universal cutoff. Calibrate it. Take a few dozen real questions, label which chunks are relevant, and look at the score distributions of relevant and irrelevant chunks. Pick the threshold that balances missing good chunks against letting bad ones through.

• Scores depend on the embedder
• Calibrate on your data

The third knob is metadata filters. Chunks carry labels like product, year, language, or which team may read them. Filtering on those fields before or during the search shrinks the candidate set to what is actually eligible. If the user asks about the twenty twenty four pricing for product B, no similarity score should ever surface a twenty nineteen page for product A.

• Filter before similarity
• Year, product, language, access

Pure similarity ranking has a flaw. If your corpus contains the same fact in five places, say a policy copied into several documents, the top five results can be near duplicates. They fill your context window with one idea, and the different chunk that holds the missing detail sits at rank six. We want relevance and diversity together.

• Top results say the same thing
• Wasted context

Maximal marginal relevance builds the result list greedily. Let q be the query, S the set already selected, and d a candidate. Score each candidate as lambda times its similarity to the query, minus one minus lambda times its highest similarity to anything already selected. The first term rewards relevance. The second punishes redundancy. Pick the best score, add it to S, and repeat.

• Reward relevance to query
• Penalize similarity to chosen

Lambda is the dial. At one, the second term vanishes and you get ordinary top k. At zero you ignore the query after the first pick and chase diversity alone, which returns irrelevant material. In practice, values between about zero point five and zero point eight are common starting points, but tune it on your own questions.

• Lambda 1: plain top-k
• Lambda 0: pure diversity
• Typical: 0.5 to 0.8

Let us run it with lambda equal to zero point seven. Three candidates. A has similarity zero point nine to the query, B zero point eight eight, C zero point seven. A and B are near duplicates, with similarity zero point nine five between them. A and C have only zero point three. Step one, nothing is selected yet, so the score is just zero point seven times the query similarity: A scores zero point six three, B zero point six one six, C zero point four nine. A wins.

• Lambda = 0.7
• A: 0.90, B: 0.88, C: 0.70

Step two, with A selected. For B, the relevance part is zero point six one six, and the redundancy penalty is zero point three times zero point nine five, which is zero point two eight five. The final score is zero point three three one. For C, relevance is zero point four nine, the penalty is zero point three times zero point three, which is zero point zero nine, giving zero point four. C beats B, even though B was more similar to the query.

• B: 0.616 - 0.285 = 0.331
• C: 0.49 - 0.09 = 0.40
• C wins

In practice you run MMR as a second step. Fetch a larger pool, say thirty candidates with their embeddings, compute the pairwise similarities between them, and pick five greedily. That is a few hundred dot products, which is negligible next to any model call. Many vector stores and frameworks offer MMR as a built in search mode.

• Fetch a larger pool first
• Run MMR over candidates
• Cost: small extra compute

Now the transformations. The most important one is fixing follow up questions. A user asks about the Pro plan, gets an answer, and then asks, and how much is it for students? Embedded alone, that sentence has no mention of the Pro plan, so retrieval has no idea what it refers to. The query depends on the conversation.

• 'What about the second one?'
• Needs the conversation

The fix is to ask a model to rewrite the follow up into a standalone query, given the recent history. The output should be something like, Pro plan pricing for students. Now it contains every term the retriever needs and nothing it does not. Note that you use the rewritten query for retrieval, but usually the original wording for the final answer.

• LLM sees history plus follow-up
• Outputs a self-contained query

In Python, this is a small function. It takes the history, the latest question, and a function that calls your model. The prompt says: rewrite the last question so it can be understood without the conversation, do not add new facts, return only the query. Use temperature zero, and if the call fails or returns empty text, fall back to the original question. Keep only the last few turns of history to control cost.

• Small prompt, low temperature
• Fall back to original

Next, multi query expansion. A single phrasing may miss a document that uses different vocabulary. So you ask a model for three to five alternative phrasings of the question, perhaps emphasizing different words or angles, and run retrieval for each. Each paraphrase lands in a slightly different spot, and together they cover more ground.

• Generate 3 to 5 paraphrases
• Cover different wordings

The same chunk will often come back from several paraphrases, so you must deduplicate by chunk ID. Two simple merge rules work well: keep the best score each chunk achieved, or count how many query variants retrieved it, since agreement is a useful signal. In code, a dictionary keyed by chunk ID does the job, and then you sort and cut to your final k.

• Same chunk, many hits
• Keep best score or count votes

The code follows that recipe. Generate n variants with the model, always add the original question to the list, search each one, and fold the results into a dictionary that keeps the maximum score per chunk ID. Finally sort by score and return the top results. Always including the original question is a cheap safeguard against a bad paraphrase.

• Generate variants
• Search each, merge by ID
• Include the original query

Here is a clever idea called HyDE, hypothetical document embeddings. Instead of embedding the question, ask the model to write a plausible answer passage, even without any knowledge of your data. Then embed that passage and search with it. An answer looks much more like a document than a question does, so its embedding tends to land closer to the real thing.

• Hypothetical Document Embeddings
• Answers look like documents

The generated passage can contain wrong facts, and that is fine, because it is only a search key. What matters is that its vocabulary and structure resemble real documents. The real retrieved chunks then supply the true facts. The risk is that for niche topics the model may invent a passage that drifts to the wrong subject, so test it. It also adds one full generation call before retrieval.

• Facts may be wrong
• Style and vocabulary matter
• Real documents ground the answer

Step back prompting goes the other way. A very specific question may match nothing, while the answer depends on a general concept. So the model first produces a broader version, for instance, turning why was I charged twice on March third into, how does billing handle duplicate charges. You retrieve for both the specific and the broad query, and give the model principles along with details.

• Ask the broader question first
• Retrieve principles, then specifics

Some questions need a chain of facts. Who manages the team that built the billing service? One document says which team built billing. Another says who manages that team. No single chunk contains both, so one search cannot find the answer. These are called multi hop questions.

• Answer needs chained facts
• One search cannot do it

Decomposition asks a model to split the question into sub questions. If they are independent, like comparing the refund rules of two plans, you retrieve for each in parallel and merge the context. If they depend on each other, as in the manager example, you run them in sequence, feeding the first answer into the next query. Sequential steps add latency, one round trip each.

• Split into sub-questions
• Sequential or parallel
• Combine at the end

Query routing decides where a query should go. You may have separate indexes for product docs and HR policies, a SQL database for numbers, or no need to retrieve at all for a greeting or a simple arithmetic question. A small classifier or a cheap model call picks the route. Routing saves cost and avoids searching the wrong corpus.

• Choose the right source
• Or skip retrieval entirely

Routers fail when routes overlap or descriptions are vague. Keep the routes few, give each a clear description and a couple of example queries, and consider a fallback: if the router is unsure, search everything. A wrong route is worse than a slower search, because the right document is never even looked at.

• Keep routes few and distinct
• Allow multi-route fallback

Every transformation has a price. A history rewrite is one short model call. Multi query is one call plus several searches, and searches run in parallel so they are cheap in time. HyDE needs a longer generation, which is slower. Step back adds a call and a second search. Decomposition can cost several sequential calls. Routing is one small call. Exact timings depend on your model and hardware, so measure.

• Rewrite: one short LLM call
• Multi-query: one call, N searches
• HyDE: one long generation
• Decomposition: several calls

The practical rule is to use the cheapest tool that fixes your observed failure. Run a small, fast model for rewriting. Apply techniques conditionally: rewrite only when there is conversation history, decompose only when the question has several parts. And measure. If a technique does not raise the share of questions where the right chunk is retrieved, drop it.

• Small model for rewriting
• Apply only when needed
• Measure recall gain

Here is MMR in plain Python. A helper computes cosine similarity. The main function keeps a list of selected indexes and a list of remaining ones. In each round it loops over the remaining candidates, computes lambda times the query similarity, subtracts one minus lambda times the maximum similarity to the selected items, and moves the best one across. Handle the first round by treating the redundancy term as zero.

• Cosine similarity helper
• Greedy selection loop

Always test with the numbers you derived by hand. Build three small vectors whose similarities match the example, run the function with lambda zero point seven and k equal to two, and confirm that it returns A and then C, not A and B. If your code disagrees with your arithmetic, one of them is wrong, and it is better to find that now.

• Expect order A, C, B
• Unit-test your logic

The first mistake is over rewriting. A rewriter can silently drop a constraint, like thirty days, or add an assumption the user never made. The query then drifts and retrieves a nicely matching but wrong chunk. Constrain the prompt, forbid new facts, log original and rewritten queries side by side, and search with both when in doubt.

• Rewriter adds or drops meaning
• Constrain the prompt
• Keep the original too

The second mistake is a top k that is too small. Teams pick three because prompts feel tidy, then wonder why answers are incomplete. Build a small set of questions with known relevant chunks and plot recall as k grows. You will usually see a curve that flattens, and the right place to fetch is near the start of the flat part.

• Right chunk at rank 7
• Measure recall at k

The third mistake is ignoring the no answer case. Retrieval always returns something, even for a question the corpus cannot answer. Those weak chunks get stuffed into the prompt, and the model, wanting to help, builds a confident answer from nothing. Use a similarity threshold, instruct the model to say it does not know when the context lacks the answer, and test with questions that have no answer.

• Always returning k chunks
• Model fabricates from weak context
• Threshold plus 'I don't know'

Let us recap. User questions are poor queries, so we reshape them. We tune top k, thresholds, and metadata filters, and use MMR to trade relevance against redundancy. We rewrite follow ups into standalone queries, expand with multiple phrasings, use HyDE and step back, decompose multi hop questions, and route to the right source. And we watch for drift, small k, and missing answers.

• Fix the query, not just the index
• Knobs: k, threshold, filters, MMR
• Transform: rewrite, multi-query, HyDE, step-back, decompose, route

We now fetch a wide, diverse set of candidates. But the order inside that set is still decided by a fairly rough similarity score. In the next episode we will fix that with reranking, using cross encoders that read the query and each chunk together.

• Retrieve wide, rerank precisely
• Cross-encoders