Self-RAG: Architecture, Self-Reflection Loops, and LangGraph Implementation
Retrieval-Augmented Generation (RAG) significantly improves Large Language Models by grounding them in external data. However, standard RAG architectures operate statically: they query vector databases unconditionally, implicitly trust semantic matches, and generate responses without verifying evidential support or topical alignment. Self-Reflective Retrieval-Augmented Generation (Self-RAG) transforms this passive flow into an active, self-correcting agentic pipeline that introspects, evaluates, and course-corrects at every stage of retrieval and synthesis.
1. Why Do We Need Self-RAG? (Failure Modes of Traditional RAG)
Traditional RAG architectures follow an unbending, feed-forward sequence:
This rigid pipeline exhibits three fundamental architectural vulnerabilities:
Structure & Concept Flow
Problem 1: Indiscriminate Retrieval
Standard RAG enforces retrieval for every user prompt, regardless of whether retrieval is necessary. If a child asks an encyclopedia-backed assistant: "How many seconds are there in a minute?", an unconstrained RAG pipeline performs a dense vector search across historical and scientific indices.
The retrieval stage might return historical perspectives on time units or colloquial definitions of a "minute." In response, the LLM dilutes its output:
"A minute typically consists of about 60 seconds, depending on historical and conversational contexts."
This exposes two flaws:
- Loss of Confidence & Precision: Forcing extraneous semantic noise into the context window causes the model to hedge simple facts using qualifiers ("typically", "about").
- Computational Overhead & Latency: Vector lookups, embedding computations, and extra context tokens waste compute and increase time-to-first-token.
Problem 2: Blind Trust in Retrieved Documents
Retrievers optimize for semantic similarity, not evidential relevance. Consider a user querying: "What causes diabetes?" The retriever may surface a high-scoring chunk stating:
"Diabetes is a chronic medical condition that affects how the body processes blood sugar."
This passage describes the effect and nature of the disease, not its cause. Because traditional RAG forces the LLM to generate an answer strictly from provided context, the model distorts the logic:
"Diabetes is caused by problems in how the body processes blood sugar."
This statement is causally flawed and factually imprecise. Traditional RAG cannot flag that the retrieved document was deficient for the user's specific causal question.
Problem 3: No Post-Generation Verification
Traditional RAG treats generation as a terminal step. Once the model outputs tokens, the system serves them directly to the end user. It performs no automated introspection to determine whether:
- The claims made in the generation are actually supported by the source chunks.
- The model hallucinated supplementary facts from its pre-training weights.
- The generated response actually resolved the core user query.
2. Core Mental Model: Self-RAG as Active Introspection
Self-RAG addresses these failure modes by turning the pipeline into a self-reflective decision engine. Rather than assuming that every step succeeded, the system explicitly self-audits by answering four foundational questions across its lifecycle:
Structure & Concept Flow
The Four Pillars of Self-Reflection
Note on Research vs. Production Implementations: In the original academic paper (Asai et al., 2023), Self-RAG was implemented by fine-tuning an open-source model (e.g., Llama-2) with specialized reflection tokens such as
[Retrieve],[IsRel],[IsSup], and[IsUse]. In modern production and LangGraph agent architectures, we achieve the exact same operational topology using structured tool calling, Pydantic schemas, and explicit conditional routing graph edges on top of enterprise-grade foundation models.
3. High-Level Architecture & End-to-End Decision Flow
Below is the complete vertical execution trace of a Self-RAG system with conditional loops for answer revision and retrieval rewriting:
Structure & Concept Flow
4. Step 1: Dynamic Retrieval Routing (decide_retrieval)
The initial gateway inspects the incoming query and evaluates whether external, proprietary, or dynamic citations are necessary, or if the model's parametric knowledge is sufficient.
Structure & Concept Flow
LangGraph Implementation: Step 1
PYTHON
from typing import TypedDict, List
from pydantic import BaseModel, Field
from langchain_openai import ChatOpenAI
from langgraph.graph import StateGraph, START, END
# --- State Definition ---
class GraphState(TypedDict):
question: str
need_retrieval: bool
docs: List[str]
answer: str
# --- Pydantic Schema for Structural Routing ---
class ShouldRetrieve(BaseModel):
should_retrieve: bool = Field(
...,
description="True if external/company documents are required to answer. False for general definitions, logic, or conversational queries."
)
llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)
# --- Node: Decide Retrieval ---
def decide_retrieval_node(state: GraphState):
structured_llm = llm.with_structured_output(ShouldRetrieve)
prompt = f"""You are an expert query classifier. Determine if external documentation is strictly needed.
Rules:
- Set should_retrieve=True if the answer requires private corporate policies, technical specs, or internal metrics.
- Set should_retrieve=False for general reasoning, common knowledge, coding syntax, or basic definitions.
- If unsure, default to True.
Query: {state['question']}"""
decision = structured_llm.invoke(prompt)
return {"need_retrieval": decision.should_retrieve}
# --- Node: Direct Generation ---
def generate_direct_node(state: GraphState):
prompt = f"""Answer the question using only your general parametric knowledge.
If you cannot answer definitively without external context, state: 'I don't know based on general knowledge.'
Question: {state['question']}"""
response = llm.invoke(prompt)
return {"answer": response.content}
5. Step 2: Chunk-Level Relevance Filtering (filter_relevance)
Standard retrievers return top- nearest neighbors based on cosine similarity, often introducing irrelevant chunks. Self-RAG passes each chunk through an individual relevance filter.
Structure & Concept Flow
LangGraph Implementation: Step 2
PYTHON
class RelevanceDecision(BaseModel):
is_relevant: bool = Field(
...,
description="True if the document contains facts directly useful for answering the query, else False."
)
class ExtendedGraphState(GraphState):
relevant_docs: List[str]
context: str
def filter_relevance_node(state: ExtendedGraphState):
structured_llm = llm.with_structured_output(RelevanceDecision)
surviving_docs = []
for doc in state.get("docs", []):
prompt = f"""You are evaluating document relevance.
Question: {state['question']}
Document:
{doc}
Does this document contain information that helps directly answer the question?"""
eval_result = structured_llm.invoke(prompt)
if eval_result.is_relevant:
surviving_docs.append(doc)
return {
"relevant_docs": surviving_docs,
"context": "\n\n".join(surviving_docs)
}
def route_after_relevance(state: ExtendedGraphState):
if len(state["relevant_docs"]) > 0:
return "generate_from_context"
return "no_relevant_docs"
6. Step 3: Evidential Grounding & Hallucination Auditing (is_supported)
When generating from retrieved context, models may draw on internal training data to inject speculative claims that were never stated in the source text.
Self-RAG grades answers across three distinct support levels:
Structure & Concept Flow
Concrete Evaluation Example
- User Query: "What are the side effects of Drug X?"
- Retrieved Evidence: "Clinical trials report mild dizziness and nausea as observed side effects."
LangGraph Implementation: Step 3 (Grounding Audit + Iterative Revision)
PYTHON
from typing import Literal
class SupportAudit(BaseModel):
support_status: Literal["fully_supported", "partially_supported", "no_support"] = Field(
...,
description="Check if every assertion in the answer is backed by the context."
)
evidence_extracted: List[str] = Field(
default_factory=list,
description="Direct quotes or sub-facts from the context backing up claims."
)
class AuditState(ExtendedGraphState):
support_status: str
evidence: List[str]
revision_retries: int
def check_hallucination_node(state: AuditState):
structured_llm = llm.with_structured_output(SupportAudit)
prompt = f"""You are a strict Hallucination and Factual Grounding Auditor.
Context:
{state['context']}
User Question: {state['question']}
Generated Answer: {state['answer']}
Audit all facts asserted in the Generated Answer against the Context. Classify as:
- fully_supported: All claims exist directly within context.
- partially_supported: Mix of true context facts and external/fabricated details.
- no_support: Claims do not exist in the context whatsoever."""
audit = structured_llm.invoke(prompt)
return {
"support_status": audit.support_status,
"evidence": audit.evidence_extracted
}
def revise_answer_node(state: AuditState):
retries = state.get("revision_retries", 0) + 1
prompt = f"""You are a Strict Editorial Auditor. The previous answer contained ungrounded or hallucinated claims.
Rewrite the answer so that it is EXCLUSIVELY supported by the Context below.
Completely remove any facts, speculation, or interpretations that do not have direct evidence in the context.
Context:
{state['context']}
Question: {state['question']}
Previous Flawed Answer: {state['answer']}"""
revised_response = llm.invoke(prompt)
return {
"answer": revised_response.content,
"revision_retries": retries
}
def route_after_hallucination_check(state: AuditState):
if state["support_status"] == "fully_supported":
return "check_usefulness"
if state.get("revision_retries", 0) >= 3:
# Prevent infinite loops; exit with partial attribution
return "check_usefulness"
return "revise_answer"
7. Step 4: Utility Auditing & Query Rewriting (is_useful)
An answer can be 100% grounded yet still fail to answer the user's question. This happens when the model generates factual summaries from the provided context that completely miss the core query.
The Grounded-Yet-Useless Failure Mode
- Question: "How many paid vacation days do Nexa AI employees receive each year?"
- Context: Chunks detailing sick leaves, remote work rules, and the HR portal.
- Model Output: "Employees at Nexa AI can take sick leaves and manage their requests via the HR portal."
- Audit:
- Grounded? Yes (Zero hallucinations; perfectly supported).
- Useful? No (The user asked for the number of paid vacation days, which remains unaddressed).
When this occurs, the model must rewrite the search query using keywords optimized for vector retrieval, and then loop back to the retrieval node.
Structure & Concept Flow
LangGraph Implementation: Step 4
PYTHON
class UsefulnessAudit(BaseModel):
is_useful: bool = Field(
...,
description="True if the response fully and directly addresses the core objective of the user question."
)
critique: str = Field(..., description="Justification for the decision.")
class QueryRewrite(BaseModel):
retrieval_query: str = Field(
...,
description="Refactored search query containing high-signal keywords optimized for dense vector retrieval."
)
class FinalSelfRAGState(AuditState):
is_useful: bool
critique: str
retrieval_query: str
query_rewrite_retries: int
def check_usefulness_node(state: FinalSelfRAGState):
structured_llm = llm.with_structured_output(UsefulnessAudit)
prompt = f"""Evaluate whether the Generated Answer actually answers the User Question.
User Question: {state['question']}
Generated Answer: {state['answer']}
Mark is_useful=True if it directly answers the core question.
Mark is_useful=False if it dodges the question, provides only tangential info, or misses the core question."""
result = structured_llm.invoke(prompt)
return {
"is_useful": result.is_useful,
"critique": result.critique
}
def rewrite_query_node(state: FinalSelfRAGState):
retries = state.get("query_rewrite_retries", 0) + 1
structured_llm = llm.with_structured_output(QueryRewrite)
prompt = f"""You are a Search Optimization Engineer.
The previous query failed to retrieve documents that could resolve the user's question.
Original Question: {state['question']}
Previous Search Query: {state.get('retrieval_query', state['question'])}
Prior Critique: {state.get('critique', 'No information found')}
Generate a new, keyword-rich search query optimized for vector retrieval over company documentation."""
rewritten = structured_llm.invoke(prompt)
return {
"retrieval_query": rewritten.retrieval_query,
"query_rewrite_retries": retries
}
def route_after_usefulness(state: FinalSelfRAGState):
if state["is_useful"]:
return END
if state.get("query_rewrite_retries", 0) >= 3:
return "terminal_no_answer"
return "rewrite_query"
8. Complete LangGraph Implementation & Graph Compilation
We now combine every conditional node, edge, and state field into an executable StateGraph.
Structure & Concept Flow
Complete Pipeline Code
PYTHON
from typing import TypedDict, List, Literal
from pydantic import BaseModel, Field
from langchain_openai import ChatOpenAI
from langgraph.graph import StateGraph, START, END
# --- Complete Pipeline State Schema ---
class SelfRAGState(TypedDict):
question: str
retrieval_query: str
need_retrieval: bool
docs: List[str]
relevant_docs: List[str]
context: str
answer: str
support_status: Literal["fully_supported", "partially_supported", "no_support"]
evidence: List[str]
revision_retries: int
is_useful: bool
critique: str
query_rewrite_retries: int
# --- Pydantic Decision Models ---
class ShouldRetrieve(BaseModel):
should_retrieve: bool = Field(..., description="True if enterprise docs needed.")
class RelevanceDecision(BaseModel):
is_relevant: bool = Field(..., description="True if chunk aids answering query.")
class SupportAudit(BaseModel):
support_status: Literal["fully_supported", "partially_supported", "no_support"]
evidence_extracted: List[str] = Field(default_factory=list)
class UsefulnessAudit(BaseModel):
is_useful: bool
critique: str
class QueryRewrite(BaseModel):
retrieval_query: str
# --- Pipeline Setup ---
llm = ChatOpenAI(model="gpt-4o", temperature=0)
# --- Node Implementations ---
def decide_retrieval(state: SelfRAGState):
structured_llm = llm.with_structured_output(ShouldRetrieve)
prompt = f"Need private docs for this query? Query: {state['question']}"
dec = structured_llm.invoke(prompt)
return {
"need_retrieval": dec.should_retrieve,
"retrieval_query": state.get("retrieval_query") or state["question"]
}
def generate_direct(state: SelfRAGState):
resp = llm.invoke(f"Answer using general knowledge: {state['question']}")
return {"answer": resp.content}
def retrieve_documents(state: SelfRAGState):
# Simulated vector retriever call using the state's retrieval query
query = state["retrieval_query"]
# doc_results = vectorstore.similarity_search(query, k=4)
# Mocking outputs for illustrative verification:
retrieved_chunks = [
f"Chunk matching {query}: Nexa AI observes 24 paid days off.",
"Irrelevant chunk about cafeteria policies."
]
return {"docs": retrieved_chunks}
def filter_relevance(state: SelfRAGState):
structured_llm = llm.with_structured_output(RelevanceDecision)
kept = []
for doc in state.get("docs", []):
prompt = f"Query: {state['question']}\nDoc: {doc}\nContains info to answer query?"
if structured_llm.invoke(prompt).is_relevant:
kept.append(doc)
return {"relevant_docs": kept, "context": "\n\n".join(kept)}
def generate_from_context(state: SelfRAGState):
prompt = f"Context:\n{state['context']}\n\nQuestion: {state['question']}\nAnswer strictly using context."
resp = llm.invoke(prompt)
return {"answer": resp.content}
def check_hallucination(state: SelfRAGState):
structured_llm = llm.with_structured_output(SupportAudit)
prompt = f"Context:\n{state['context']}\nQuestion: {state['question']}\nAnswer: {state['answer']}\nAudit facts."
audit = structured_llm.invoke(prompt)
return {"support_status": audit.support_status, "evidence": audit.evidence_extracted}
def revise_answer(state: SelfRAGState):
prompt = f"Context:\n{state['context']}\nFlawed Answer: {state['answer']}\nRewrite to remove unsupported claims."
resp = llm.invoke(prompt)
return {
"answer": resp.content,
"revision_retries": state.get("revision_retries", 0) + 1
}
def check_usefulness(state: SelfRAGState):
structured_llm = llm.with_structured_output(UsefulnessAudit)
prompt = f"Question: {state['question']}\nAnswer: {state['answer']}\nDoes this resolve the question?"
audit = structured_llm.invoke(prompt)
return {"is_useful": audit.is_useful, "critique": audit.critique}
def rewrite_query(state: SelfRAGState):
structured_llm = llm.with_structured_output(QueryRewrite)
prompt = f"Original: {state['question']}\nPrior Query: {state['retrieval_query']}\nCritique: {state['critique']}\nOptimize query keywords for vector DB."
res = structured_llm.invoke(prompt)
return {
"retrieval_query": res.retrieval_query,
"query_rewrite_retries": state.get("query_rewrite_retries", 0) + 1
}
def no_relevant_docs_fallback(state: SelfRAGState):
return {"answer": "No relevant documentation could be located to resolve your inquiry."}
# --- Routing Logic Functions ---
def route_decide_retrieval(state: SelfRAGState):
return "retrieve_documents" if state["need_retrieval"] else "generate_direct"
def route_relevance(state: SelfRAGState):
return "generate_from_context" if len(state["relevant_docs"]) > 0 else "no_relevant_docs_fallback"
def route_hallucination(state: SelfRAGState):
if state["support_status"] == "fully_supported":
return "check_usefulness"
if state.get("revision_retries", 0) >= 3:
return "check_usefulness"
return "revise_answer"
def route_usefulness(state: SelfRAGState):
if state["is_useful"]:
return END
if state.get("query_rewrite_retries", 0) >= 3:
return "no_relevant_docs_fallback"
return "rewrite_query"
# --- Graph Assembly ---
workflow = StateGraph(SelfRAGState)
workflow.add_node("decide_retrieval", decide_retrieval)
workflow.add_node("generate_direct", generate_direct)
workflow.add_node("retrieve_documents", retrieve_documents)
workflow.add_node("filter_relevance", filter_relevance)
workflow.add_node("generate_from_context", generate_from_context)
workflow.add_node("no_relevant_docs_fallback", no_relevant_docs_fallback)
workflow.add_node("check_hallucination", check_hallucination)
workflow.add_node("revise_answer", revise_answer)
workflow.add_node("check_usefulness", check_usefulness)
workflow.add_node("rewrite_query", rewrite_query)
# Defining Edges
workflow.add_edge(START, "decide_retrieval")
workflow.add_conditional_edges(
"decide_retrieval",
route_decide_retrieval,
{"retrieve_documents": "retrieve_documents", "generate_direct": "generate_direct"}
)
workflow.add_edge("generate_direct", END)
workflow.add_edge("retrieve_documents", "filter_relevance")
workflow.add_conditional_edges(
"filter_relevance",
route_relevance,
{"generate_from_context": "generate_from_context", "no_relevant_docs_fallback": "no_relevant_docs_fallback"}
)
workflow.add_edge("no_relevant_docs_fallback", END)
workflow.add_edge("generate_from_context", "check_hallucination")
workflow.add_conditional_edges(
"check_hallucination",
route_hallucination,
{"check_usefulness": "check_usefulness", "revise_answer": "revise_answer"}
)
workflow.add_edge("revise_answer", "check_hallucination")
workflow.add_conditional_edges(
"check_usefulness",
route_usefulness,
{"__end__": END, "rewrite_query": "rewrite_query", "no_relevant_docs_fallback": "no_relevant_docs_fallback"}
)
workflow.add_edge("rewrite_query", "retrieve_documents")
app = workflow.compile()
9. Production Considerations & Engineering Pitfalls
When moving Self-RAG from prototyping to production, watch out for these practical challenges:
Structure & Concept Flow
1. Latency Accumulation
- The Problem: In the worst-case scenario (retrieval document filter generation 3 revisions query rewrite re-retrieval generation), an end-user might wait for 8 to 12 serial LLM invocations.
- Production Fix:
- Use smaller, high-throughput models (e.g.,
gpt-4o-mini,claude-3-haiku) for the routing, relevance, and audit nodes. - Reserve larger models (e.g.,
gpt-4o,claude-3-5-sonnet) exclusively for the final contextual generation node. - Run the
filter_relevanceevaluations across chunks in parallel usingasyncio.gather.
- Use smaller, high-throughput models (e.g.,
2. Guardrails Against Infinite Loops
- The Problem: A low-confidence or conflicting vector index can cause the agent to bounce endlessly between
revise_answercheck_hallucinationorrewrite_queryretrieve_documents. - Production Fix:
- Always track loop counters in the graph state (e.g.,
revision_retries,query_rewrite_retries). - Enforce hard bounds (e.g., retries) before degrading gracefully into an explicit failure state (
no_answer_found).
- Always track loop counters in the graph state (e.g.,
3. Relevance Thresholding
- The Problem: If the
filter_relevanceprompt is overly aggressive, it may discard chunks that contain indirect or distributed context, leaving the model with an empty context window. - Production Fix:
- Instruct the model that if an extracted chunk provides partial or background signal for the query, it should be marked as relevant (
is_relevant = True). - Treat relevance as high-recall, low-precision filtering, letting the downstream hallucination auditor handle factual alignment.
- Instruct the model that if an extracted chunk provides partial or background signal for the query, it should be marked as relevant (
10. Summary and Key Takeaways
- Traditional RAG Fails Statically: Conventional RAG pipelines suffer from indiscriminate retrieval (wasting compute on basic queries), blind trust in semantically matched chunks (injecting irrelevance), and an absence of output verification.
- Self-RAG Enables Dynamic Introspection: Self-RAG introduces four self-reflection checkpoints:
- Should we retrieve?
- Are the retrieved documents relevant?
- Is the generated answer grounded in the evidence?
- Does the answer actually satisfy the user's intent?
- Tri-State Hallucination Classification: Generated claims are categorized as
fully_supported,partially_supported, orno_support. Only fully supported answers reach the user; ungrounded generations are routed through a revision loop. - Decoupling Grounding from Utility: An answer can be completely grounded in context while still failing to answer the user's question. Self-RAG explicitly separates hallucination auditing from utility auditing.
- Query Rewriting Closes the Loop: When context proves unhelpful, Self-RAG refactors the prompt into an optimized search query and loops back to retrieval, rather than surfacing irrelevant facts.
- Clean Implementation with Modern Frameworks: While the original research relied on specialized fine-tuned reflection tokens, modern production architectures can implement the complete Self-RAG pattern using LangGraph, Pydantic schemas, and structured outputs.
- Engineered Loop Termination: Production implementations must enforce strict retry limits on revision and retrieval loops to balance reflection quality against latency and token costs.