Skip to main content

Query rewriting through self-inflection

Retrieval can fail silently. Hybrid search always returns something in the top-k, and a reranker always produces an ordering. Neither step guarantees that the retained chunks actually contain what is needed to answer the question. If you prompt the LLM anyway, you often get a fluent answer that is under-grounded — or confidently wrong.

RAG Me Up's rewrite loop is an optional self-inflection / LLM-as-a-judge guardrail for that failure mode. After documents are retrieved (and optionally reranked), another LLM call asks: can these documents answer the question? If the judge says no, the question is rewritten once and retrieval is repeated.

"Self-inflection" here means the system critiques its own retrieval result and adapts the query — not the final answer — before generation.

Why only once?

Self-correction loops are easy to turn into latency sinks or infinite cycles. RAG Me Up allows at most one rewrite. That recovers many weak first-pass queries (missing synonyms, awkward phrasing, underspecified entities) without thrashing.

The loop is also disabled when HyDE is on. HyDE already replaced the user question with a hypothetical document; judging and rewriting on top of that is the wrong abstraction.

The two-step judge

1. Can the documents answer?

(response, _) = self.llm.generate_response(
os.getenv("rewrite_query_instruction").format(context=self.format_documents(documents)),
os.getenv("rewrite_query_question").format(question=prompt),
[]
)

rewrite_query_instruction is a system prompt that must include the retrieved {context} and force a yes/no decision (plus a motivation). rewrite_query_question wraps the user {question}.

If the reply starts with no, the pipeline continues to rewrite. Anything else is treated as "documents are good enough" (fail open toward answering).

2. Rewrite the question for better retrieval

(new_prompt, _) = self.llm.generate_response(
None,
os.getenv("rewrite_query_prompt").format(
question=prompt,
motivation=f"Can I find the answer in the documents: {response}"
),
[]
)
rewritten = new_prompt
documents = self.handle_documents(new_prompt, prompt_embedding, datasets)

The judge's motivation is fed into the rewrite prompt so the model knows why the first retrieval failed (e.g. "documents discuss maternity leave but not partner leave"). The rewritten question then runs through the same handle_documents path (retrieve + optional rerank).

The UI can surface rewritten so users see that their question was adapted — useful for trust and for debugging.

Configuration

VariableRole
use_rewrite_loopEnable/disable the judge + single rewrite
rewrite_query_instructionSystem prompt with {context}; force yes/no
rewrite_query_questionUser message with {question}
rewrite_query_promptRewrite instruction with {question} and {motivation}

Designing the prompts

  • Force a strict answer format (yes / no as first token) so startswith("no") is reliable.
  • Ask for a short motivation even on "yes" — invaluable in DEBUG logs when hunting false negatives.
  • Tell the judge to say no when documents are only tangentially related, not merely when they are empty.
  • Tell the rewriter to output only the rewritten question — no "Sure, here's a better phrasing:".

Educational takeaway

LLM-as-a-judge patterns show up in eval, moderation, and routing. Here the judge scores retrieval adequacy before generation. Put cheap checks as early as you can; spend expensive answer generation only on contexts that deserve it.