Query rewriting agentic RAG systems close the gap between how users phrase questions and how knowledge bases are indexed. A raw query like “fix the auth bug” rarely matches the doc titled “Troubleshooting OAuth token refresh failures.” This guide walks through adding a rewriting stage to an existing agentic loop with runnable Python and concrete verification.
Step 1: Audit your existing agentic RAG retrieval path
Most agentic RAG implementations start with a single retrieve-then-generate call. The agent sends the user message straight to a vector store and hopes the top-k chunks answer the question.
def agent_loop(user_msg: str):
docs = retrieve(user_msg) # single-string query
return llm_generate(user_msg, docs)Trace where retrieve is invoked. If it is called once per user turn with no transformation, you have a lexical mismatch risk. In multi-step agents, the same pattern often hides inside a tool call. Extract that call into a named function so you can intercept it later.
Step 2: Specify the rewriting contract
The rewriter must output structured, deterministic data. Free-text rewrites are undebuggable. Define a minimal JSON schema: a list of strings, each a standalone search query.
{
"queries": ["OAuth token refresh failure troubleshooting", "auth bug fix guide"]
}Write a system prompt that enforces decomposition, expansion, and entity normalization. For example: “Rewrite the user query into 1–3 precise retrieval queries. Expand abbreviations, add synonyms, and split multi-part questions. Output only JSON matching {"queries": [string]}.”
Keep the count low. Three queries is enough to cover intent without blowing up retrieval latency or context size.
Step 3: Implement the rewriter with a dedicated model
Use a smaller, cheaper chat model for rewriting and reserve the heavy model for synthesis. n4n.ai exposes an OpenAI-compatible endpoint covering 240+ models with automatic fallback when a provider is rate-limited, so you can pin a fast model and not worry about transient 429s.
import json
from openai import OpenAI
client = OpenAI(base_url="https://api.n4n.ai/v1", api_key="YOUR_KEY")
def rewrite_query(user_msg: str) -> list[str]:
resp = client.chat.completions.create(
model="gpt-4o-mini",
messages=[
{"role": "system", "content": "Rewrite the user query into 1-3 precise search queries. Return JSON: {'queries': [str]}"},
{"role": "user", "content": user_msg}
],
response_format={"type": "json_object"}
)
return json.loads(resp.choices[0].message.content)["queries"]The response_format guard forces valid JSON on compliant models. If your model does not support it, parse with a strict validator and fall back to [user_msg] on error.
Step 4: Fan out retrieval across rewritten queries
Run each query through your existing retrieve in parallel. Deduplicate by document id and preserve source query for debugging.
import concurrent.futures
from dataclasses import dataclass
@dataclass
class Doc:
id: str
text: str
score: float
def retrieve_many(queries: list[str]) -> list[Doc]:
with concurrent.futures.ThreadPoolExecutor() as ex:
results = ex.map(retrieve, queries)
seen = {}
for docs in results:
for d in docs:
if d.id not in seen or d.score > seen[d.id].score:
seen[d.id] = d
return sorted(seen.values(), key=lambda x: x.score, reverse=True)Thread pool overhead is negligible compared to network latency on vector search. Cap concurrency at 5 to avoid hammering the index.
Step 5: Merge and rank retrieved context
Rewritten queries surface overlapping chunks. Truncate to the agent’s context budget after global ranking. Do not concatenate per-query top-k blindly; that biases toward the first query.
def build_context(docs: list[Doc], max_tokens: int = 3000) -> str:
out, used = [], 0
for d in docs:
if used + len(d.text) // 4 > max_tokens:
break
out.append(d.text)
used += len(d.text) // 4
return "\n\n".join(out)A simple token heuristic (len//4) avoids importing a tokenizer in the hot path. Swap in a real tokenizer if you need exact limits.
Step 6: Modify the agent loop and preserve state
Drop the rewriter into the existing loop. If the agent is ReAct-style, call rewrite_query at the start of each retrieve tool invocation, not just the first turn.
def agent_loop_v2(user_msg: str):
queries = rewrite_query(user_msg)
docs = retrieve_many(queries)
context = build_context(docs)
return llm_generate(user_msg, context)For stateful agents, attach queries to the trace so later steps can see what was searched. This makes the query rewriting agentic RAG behavior auditable when the agent takes corrective actions.
Step 7: Add observability and cache hints
Log the original query, rewritten queries, and hit counts per query. Per-token usage metering (available on gateways like n4n.ai) lets you attribute cost to the rewrite step versus synthesis. Forward provider cache-control hints on the rewrite call if your gateway honors them; rewriter prompts are static and cache well.
resp = client.chat.completions.create(
model="gpt-4o-mini",
messages=[...],
response_format={"type": "json_object"},
extra_headers={"cache-control": "max-age=3600"} # hint, honored by some providers
)Do not cache user-specific queries. Only cache the system prompt and static few-shot examples.
Step 8: Verify success with offline and online checks
Verification is where most teams cut corners. Write a golden set of 20–50 real queries with known relevant doc ids. Assert the rewritten fan-out retrieves at least one gold doc.
def test_rewrite_recall():
gold = [("auth bug", "oauth_refresh_doc")]
for q, doc_id in gold:
queries = rewrite_query(q)
docs = retrieve_many(queries)
assert doc_id in [d.id for d in docs]Run this in CI against a frozen index. Online, track retrieval hit rate (did the agent cite a retrieved chunk?) and end-task accuracy. A successful query rewriting agentic RAG rollout shows higher recall@5 on the golden set and fewer “I couldn’t find” agent failures, without measurable latency regression beyond the rewrite call.
If latency spikes, move the rewriter to a smaller model or batch rewrite multiple pending agent steps. The architecture above isolates that change to one function.