Self-RAG
On this page
Self-Reflection in RAG
Corrective RAG controls what goes into the prompt, but not what comes out. Even perfect context doesn’t guarantee a good answer: the model might add details found nowhere in the documents, or produce text that’s accurate but dodges the actual question. Self-RAG (Self-Reflective RAG) closes this gap by adding two checks after generation. The first examines factual grounding—does the answer actually rely on the provided documents? The second examines usefulness—does it truly address what was asked? If either check fails, the system loops and retries, making it a self-correction mechanism rather than a simple filter.
The flow therefore extends Corrective RAG. You still have retrieval followed by document grading, then generation. Two new nodes come next: the hallucination check, which compares the answer against the documents, and the answer check, which compares it against the question. If a check fails, the system either returns to generation or fetches fresh material from the web, depending on which verification failed.
Checking Factual Grounding
The first checker receives two inputs: the concatenated documents and the generated answer. Its job isn’t to judge whether the answer is true in absolute terms, but whether it’s entirely grounded in what it was given. This distinction matters: an accurate statement absent from the documents must be rejected because it comes from the model’s training memory, not your corpus—and there’s no guarantee the next statement will be accurate. The output format matches the grader: a binary JSON object at temperature zero.
import json
from mistralai.client import Mistral
client = Mistral(api_key=api_key)
def check_hallucination(documents, generation):
"""Check whether the generation is grounded in the documents."""
docs_text = "\n\n".join([d["text"] for d in documents])
prompt = f"""You are a factual grounding evaluator. Given a set of
documents and a generated answer, determine whether the answer is entirely
grounded in the information contained in the documents.
Respond ONLY with a JSON object: {{"grounded": "yes"}} or {{"grounded": "no"}}
Documents:
{docs_text}
Generated answer:
{generation}
JSON evaluation:"""
response = client.chat.complete(
model="mistral-large-latest",
messages=[{"role": "user", "content": prompt}],
temperature=0,
response_format={"type": "json_object"},
)
result = json.loads(response.choices[0].message.content)
return result.get("grounded", "no") == "yes"Checking Usefulness
The second checker sees neither the documents nor the retrieval results—only the question and the answer. This narrow field of view is deliberate because it isolates a flaw the first check lets through. Imagine someone asks how many dimensions an embeddings model produces, and the answer describes at length what embeddings are used for without ever giving the number. Everything is grounded in the documents, so the first checker says “yes,” and the user walks away without their information. The second checker says “no” and triggers another attempt.
def check_answer_usefulness(question, generation):
"""Check whether the answer actually addresses the question."""
prompt = f"""You are an answer quality evaluator. Given a question
and an answer, determine whether the answer usefully and completely addresses the question.
Respond ONLY with a JSON object: {{"useful": "yes"}} or {{"useful": "no"}}
Question: {question}
Answer: {generation}
JSON evaluation:"""
response = client.chat.complete(
model="mistral-large-latest",
messages=[{"role": "user", "content": prompt}],
temperature=0,
response_format={"type": "json_object"},
)
result = json.loads(response.choices[0].message.content)
return result.get("useful", "no") == "yes"The Self-Correction Loop
The SelfRAG class reuses the familiar retrieval and grading steps, then wraps them in a retry loop bounded by max_retries. The loop logic is subtler than it appears because the two failure types receive different treatment. A detected hallucination simply regenerates with the same context: the problem lies with the model, which overreached, and a fresh draw often suffices. An answer judged useless, however, triggers an additional web search before retrying: here the problem lies with the context, which is too thin to support a complete answer, and regenerating without adding anything would yield the same result.
Note the fallback exit after the loop. When all attempts are exhausted, the system doesn’t raise an exception: it returns the last generation along with its grounded and useful flags. This is a deliberate design choice—better to return an imperfect answer labeled as such, which your interface can flag for the user, than to throw a raw error.
import numpy as np
class SelfRAG:
"""Self-RAG pipeline with self-correction."""
def __init__(self, chunks, embeddings, index, max_retries=2):
self.chunks = chunks
self.embeddings = embeddings
self.index = index
self.max_retries = max_retries
def retrieve(self, question, k=4):
"""Standard retrieval."""
response = client.embeddings.create(
model="mistral-embed", inputs=[question]
)
query_vec = np.array([response.data[0].embedding], dtype=np.float32)
distances, indices = self.index.search(query_vec, k)
return [
{"text": self.chunks[idx]["text"], "source": self.chunks[idx].get("source", "")}
for dist, idx in zip(distances[0], indices[0]) if idx >= 0
]
def grade_documents(self, question, documents):
"""Evaluate document relevance."""
relevant = []
for doc in documents:
if grade_document(question, doc["text"]):
relevant.append(doc)
return relevant, len(relevant) < len(documents)
def generate(self, question, documents):
"""Generate an answer."""
context = "\n\n---\n\n".join([d["text"] for d in documents])
prompt = f"Context:\n{context}\n\nQuestion: {question}"
response = client.chat.complete(
model="mistral-large-latest",
messages=[{"role": "user", "content": prompt}],
temperature=0.2,
)
return response.choices[0].message.content
def ask(self, question, k=4):
"""Complete Self-RAG pipeline with self-correction."""
print(f"\n=== Self-RAG: {question[:60]}... ===")
# 1. Retrieval
print("1. Retrieval...")
docs = self.retrieve(question, k=k)
# 2. Grading
print("2. Grading documents...")
relevant_docs, need_web = self.grade_documents(question, docs)
if need_web:
print(" -> Fallback web search")
web_docs = web_search_fallback(question)
relevant_docs.extend(web_docs)
# 3. Generation loop with self-correction
for attempt in range(self.max_retries + 1):
print(f"3. Generation (attempt {attempt + 1})...")
generation = self.generate(question, relevant_docs)
# 4. Hallucination check
print("4. Checking hallucinations...")
is_grounded = check_hallucination(relevant_docs, generation)
if not is_grounded:
print(" -> HALLUCINATION detected, retrying...")
continue
# 5. Usefulness check
print("5. Checking usefulness...")
is_useful = check_answer_usefulness(question, generation)
if not is_useful:
print(" -> Answer NOT USEFUL, fallback web search...")
web_docs = web_search_fallback(question)
relevant_docs.extend(web_docs)
continue
# Success
print(" -> Answer validated!")
return {
"answer": generation,
"attempts": attempt + 1,
"grounded": True,
"useful": True,
}
# If all attempts fail
print(" -> Maximum attempts reached")
return {
"answer": generation,
"attempts": self.max_retries + 1,
"grounded": is_grounded,
"useful": is_useful,
}
# Usage
self_rag = SelfRAG(chunks, embeddings, index, max_retries=2)
result = self_rag.ask("What are the advantages of RAG over fine-tuning?")
print(f"\nAnswer: {result['answer']}")
print(f"Attempts: {result['attempts']}")
print(f"Grounded: {result['grounded']}, Useful: {result['useful']}")The Cost of Reliability
Each guardrail adds another LLM call, and the total adds up quickly. The function below makes this explicit by breaking down calls per step: one grading call per document, then generation, grounding check, and usefulness check, each repeated on every retry. With four documents and a single retry, you reach ten calls for a single question—where basic RAG required two. Run this estimate with your own values before setting k and max_retries: this calculation tells you whether your budget and latency requirements support Self-RAG across all traffic, or only for the most sensitive questions.
def estimate_self_rag_cost(n_docs=4, n_retries=1):
"""Estimate the number of LLM calls for a Self-RAG question."""
calls = {
"grading": n_docs, # 1 call per document
"generation": 1 + n_retries, # 1 initial + retries
"hallucination_check": 1 + n_retries,
"answer_check": 1 + n_retries,
}
total = sum(calls.values())
print(f"Estimated LLM calls: {total}")
for step, count in calls.items():
print(f" {step}: {count}")
return total
estimate_self_rag_cost(n_docs=4, n_retries=1)
# Estimated LLM calls: 10Key Takeaways
- Self-RAG adds two post-generation checks: hallucinations and usefulness
- If the answer isn’t grounded in the documents, the system regenerates
- If the answer isn’t useful, the system performs a web search and retries
- The number of retries must be limited to control latency
- Self-RAG is more expensive (more LLM calls) but more reliable