Query Rewriting and Multi-Query Retrieval

Users rarely type the perfect search query. They write vague, short or chatty questions, and a retriever matches those words against your documents. Query rewriting asks an LLM to turn the user’s question into a better search query. Multi-query retrieval goes one step further: it asks for several different versions of the question, searches with each one, and merges the results. Both techniques raise the chance that the right chunk reaches the final prompt.

Overview: how it works

A retriever has no real understanding of intent. A vector store compares the embedding of the query with the embeddings of chunks. A keyword index compares words. In both cases the wording of the query matters a lot. Look at the question “How do I stop paying?”. Your help article may say “cancel a subscription”. The two sentences share almost no words, and the embedding may place them far apart too.

Query transformation fixes this gap by adding one small LLM call before retrieval. The flow looks like this:

1. The user sends a question.
2. A prompt asks an LLM to rewrite it, or to write N alternative versions.
3. The retriever runs once per generated query.
4. The results are merged and duplicates are removed.
5. The merged documents go to the answer-generation prompt as usual.

Query rewriting

Rewriting produces one improved query. Typical jobs for the rewriter are:

• Replace slang or vague words with domain terms (“stop paying” becomes “cancel subscription”).
• Remove filler such as “hey, can you please tell me”.
• Resolve references in a chat. If the user says “and what about its price?”, the rewriter uses the chat history to produce “What is the price of the Pro plan?”. This is called a standalone question.

Multi-query retrieval

Multi-query produces several queries, for example three to five. Each one looks at the problem from a different angle: different vocabulary, a more general view, a more specific view. Because each query lands in a slightly different place in the vector space, the union of results has better recall (the share of relevant documents that you actually find). The price is extra latency and cost: one more LLM call, plus N searches.

Related ideas

• Step-back prompting: ask a broader question first, then retrieve for it as well.
• HyDE (hypothetical document embeddings): let the LLM write a fake answer and search with that text.
• Reciprocal rank fusion (RRF): a way to merge ranked lists by rank position instead of just removing duplicates.

Syntax and key concepts

With LCEL (LangChain Expression Language, the | operator) the general shape of a multi-query chain is:

prompt | llm | split_into_queries | retriever.map() | merge_and_dedupe

LangChain also ships a ready-made class, MultiQueryRetriever (import path langchain.retrievers.multi_query), which wraps a base retriever and an LLM and does the same job. It is shown in Example 3. Check the current LangChain documentation, because import paths in the ecosystem can move between releases.

Examples

The first two examples run fully offline. We use a tiny keyword retriever and a fake “LLM” that returns fixed text. This lets you see the exact data flow without an API key. The knowledge base is five short help-centre sentences.

Example 1: a single rewrite

import re
from langchain_core.documents import Document
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnableLambda

STOP = {'a', 'an', 'and', 'are', 'after', 'be', 'can', 'choose', 'do', 'from',
        'how', 'i', 'is', 'my', 'of', 'on', 'open', 'the', 'to', 'uses', 'your'}

DOCS = [
    Document(page_content='Reset your password from the account settings page.'),
    Document(page_content='Refunds are issued within 5 business days after approval.'),
    Document(page_content='To cancel a subscription, open billing and choose cancel plan.'),
    Document(page_content='Invoices can be downloaded from the billing page.'),
    Document(page_content='Two factor login uses an authenticator app.'),
]


def tokens(text):
    return {w for w in re.findall(r'[a-z0-9]+', text.lower()) if w not in STOP}


def keyword_search(query, k=2):
    q = tokens(query)
    scored = [(len(q & tokens(d.page_content)), d) for d in DOCS]
    scored = [s for s in scored if s[0] > 0]
    scored.sort(key=lambda s: s[0], reverse=True)
    return [d for _, d in scored[:k]]


retriever = RunnableLambda(keyword_search)

rewrite_prompt = ChatPromptTemplate.from_template(
    'Rewrite the question as a short search query using help-centre terms.\nQuestion: {question}'
)
fake_rewriter = RunnableLambda(lambda _: 'cancel subscription')
rewrite_chain = rewrite_prompt | fake_rewriter

question = 'How do I stop paying?'
print('Original:', [d.page_content for d in retriever.invoke(question)])
better = rewrite_chain.invoke({'question': question})
print('Rewritten:', better)
print('Result:', [d.page_content for d in retriever.invoke(better)])

Output:

Original: []
Rewritten: cancel subscription
Result: ['To cancel a subscription, open billing and choose cancel plan.']

The original question shares no meaningful word with any document, so nothing is found. After the rewrite, the retriever finds the cancellation article. In a real system fake_rewriter is a chat model, and the rewritten text would vary a little between runs.

Example 2: multi-query with merge and deduplication

import re
from langchain_core.documents import Document
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnableLambda

STOP = {'a', 'an', 'and', 'are', 'after', 'be', 'can', 'choose', 'do', 'from',
        'how', 'i', 'is', 'my', 'of', 'on', 'open', 'the', 'to', 'uses', 'your'}

DOCS = [
    Document(page_content='Reset your password from the account settings page.'),
    Document(page_content='Refunds are issued within 5 business days after approval.'),
    Document(page_content='To cancel a subscription, open billing and choose cancel plan.'),
    Document(page_content='Invoices can be downloaded from the billing page.'),
    Document(page_content='Two factor login uses an authenticator app.'),
]


def tokens(text):
    return {w for w in re.findall(r'[a-z0-9]+', text.lower()) if w not in STOP}


def keyword_search(query, k=2):
    q = tokens(query)
    scored = [(len(q & tokens(d.page_content)), d) for d in DOCS]
    scored = [s for s in scored if s[0] > 0]
    scored.sort(key=lambda s: s[0], reverse=True)
    return [d for _, d in scored[:k]]


retriever = RunnableLambda(keyword_search)

multi_prompt = ChatPromptTemplate.from_template(
    'Write 3 different search queries, one per line, for this question: {question}'
)
fake_llm = RunnableLambda(lambda _: 'cancel subscription\nbilling invoices\nrefunds days')


def split_lines(text):
    return [line.strip() for line in text.splitlines() if line.strip()]


def unique_union(batches):
    seen = set()
    merged = []
    for batch in batches:
        for doc in batch:
            if doc.page_content not in seen:
                seen.add(doc.page_content)
                merged.append(doc)
    return merged


multi_query = (
    multi_prompt
    | fake_llm
    | RunnableLambda(split_lines)
    | retriever.map()
    | RunnableLambda(unique_union)
)

for doc in multi_query.invoke({'question': 'How do I stop paying?'}):
    print('-', doc.page_content)

Output:

- To cancel a subscription, open billing and choose cancel plan.
- Invoices can be downloaded from the billing page.
- Refunds are issued within 5 business days after approval.

Three queries run. The first finds the cancellation article. The second finds the invoice article first and the cancellation article again, because both contain “billing”. The third finds the refund article. unique_union drops the repeated cancellation article and keeps the order in which documents first appeared.

Example 3: the production version

import logging
import os

from langchain.retrievers.multi_query import MultiQueryRetriever
from langchain_openai import ChatOpenAI

# Reads OPENAI_API_KEY from the environment; model name is also configurable.
llm = ChatOpenAI(model=os.environ['CHAT_MODEL'], temperature=0)

# 'vectorstore' is any vector store you built earlier in the course.
base_retriever = vectorstore.as_retriever(search_kwargs={'k': 4})

retriever = MultiQueryRetriever.from_llm(retriever=base_retriever, llm=llm)

logging.basicConfig()
logging.getLogger('langchain.retrievers.multi_query').setLevel(logging.INFO)

docs = retriever.invoke('How do I stop paying?')
print(len(docs))

This needs a provider key and a populated vector store, so its output is illustrative: with logging on, you see the generated queries printed, and docs holds the unique documents from all of them. Use a model name that is current for your provider; model names change often, so check the provider’s documentation. Setting temperature=0 keeps the generated queries more stable.

Under the hood: step by step

1. Prompt formatting. The ChatPromptTemplate fills {question} and produces a prompt value.
2. LLM call. The model returns text. In MultiQueryRetriever an output parser splits it by newline into a list of queries.
3. Fan-out. retriever.map() applies the retriever to every query. LangChain can run these in parallel, which keeps latency close to one search.
4. Fan-in. The list of lists is flattened. MultiQueryRetriever takes the unique union of the documents.
5. Generation. The merged documents are formatted into context for the answer prompt.

Note that the original question is not always searched. Many implementations add it to the query list so that a bad rewrite cannot make things worse. It is a good habit to include it yourself.

Common Mistakes

1. Merging without deduplication

batches = retriever.map().invoke(queries)
context_docs = [doc for batch in batches for doc in batch]  # wrong

Different queries often return the same chunk. Without deduplication the prompt contains repeated text, which wastes tokens and can bias the model towards that chunk. Fix: use a set of seen keys (like unique_union above), or a stable document ID.

2. A vague rewrite prompt

prompt = ChatPromptTemplate.from_template('Improve this: {question}')  # wrong

The model may answer the question, add a long explanation, or change its meaning. Be strict about the output format and the goal:

prompt = ChatPromptTemplate.from_template(
    'Rewrite the question as one short search query. '
    'Keep all names, numbers and constraints. Output only the query.\n'
    'Question: {question}'
)

3. Losing important constraints

A rewrite of “refund policy for orders in 2023” into “refund policy” drops the year. Tell the model to keep dates, product names and numbers, and test with questions that contain such details.

4. Rewriting every single query

Short, precise queries (an error code, a product ID) often work best as they are. Always rewriting adds cost and can damage them. Consider skipping the rewrite for such inputs.

Best Practices

• Keep the original question in the query set, next to the generated ones.
• Generate a small number of queries (three to five). More queries mean more cost, and the extra ones are usually near-duplicates.
• Use temperature=0 and a clear output format, one query per line, with no numbering.
• Use a small, fast model for rewriting. The task is simple and sits on the critical path of every request.
• Deduplicate by a stable key such as a chunk ID or source plus offset, not only by text.
• For chat apps, include the recent chat history in the rewrite prompt so follow-up questions become standalone.
• Log the generated queries. When an answer is wrong, the logs tell you if retrieval or rewriting was at fault.
• Measure the effect. Build a small set of test questions with known correct documents and compare recall with and without rewriting.
• Add timeouts and a fallback. If the rewrite call fails, search with the original question.
• If you need ranking, consider reciprocal rank fusion instead of a plain union.

Practice Exercises

1. Add the original query. Change Example 2 so the original question is searched together with the three generated queries. What does the final list look like? Hint: build the list in split_lines, or add a step that appends the question.
2. Limit the output. Modify unique_union to accept a max_docs value and stop after that many documents. Why might you want this before sending context to an LLM?
3. Standalone question. Write a ChatPromptTemplate with the variables {history} and {question} that rewrites a follow-up such as “and what about its price?” into a standalone question. Test it with a fake model first, then a real one.

Summary

• Retrievers are sensitive to wording, so raw user questions often miss the right chunks.
• Query rewriting produces one better query. Multi-query produces several and merges the results to improve recall.
• In LCEL the pattern is prompt, LLM, split, retriever.map(), then merge and deduplicate.
• MultiQueryRetriever packages this pattern; check the current docs for its import path.
• Strict prompts, a small query count, deduplication, logging and a fallback to the original question make the technique safe in production.
• The cost is extra latency and tokens, so measure whether recall really improves for your data.