Corrective RAG (CRAG) in Python with LangGraph: Building Self-Reflective Retrieval Pipelines with Document Grading, Fallback Web Search, and Hallucination Gates

python3 -m venv .venv
source .venv/bin/activate
pip install langchain langchain-openai langchain-community langgraph chromadb tavily-python pydantic tiktoken

export OPENAI_API_KEY="your-openai-api-key"
export TAVILY_API_KEY="your-tavily-api-key"

from typing import List, Dict, Any, Literal
from typing_extensions import TypedDict
from pydantic import BaseModel, Field
from langchain_core.documents import Document

class GraphState(TypedDict):
 """
 Represents the operational state of our Corrective RAG graph.

 Attributes:
 question: The user input query.
 generation: The synthesized response generated by the LLM.
 web_search: Boolean flag indicating if web search augmentation is required.
 documents: List of retrieved and filtered Document objects.
 retry_count: Circuit breaker to prevent infinite regeneration loops.
 """
 question: str
 generation: str
 web_search: bool
 documents: List[Document]
 retry_count: int

# Pydantic Schemas for Structured LLM Evaluation
class GradeDocuments(BaseModel):
 """Binary score for relevance check on retrieved documents."""
 binary_score: Literal["yes", "no"] = Field(
 description="Documents are relevant to the question, 'yes' or 'no'"
 )

class GradeHallucinations(BaseModel):
 """Binary score for groundedness verification."""
 binary_score: Literal["yes", "no"] = Field(
 description="Answer is grounded in the facts, 'yes' or 'no'"
 )

class GradeAnswer(BaseModel):
 """Binary score for query-answer completeness."""
 binary_score: Literal["yes", "no"] = Field(
 description="Answer addresses the user question, 'yes' or 'no'"
 )

from langchain_openai import ChatOpenAI, OpenAIEmbeddings
from langchain_community.vectorstores import Chroma
from langchain_core.prompts import ChatPromptTemplate

# Initialize LLMs
evaluator_llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)
generator_llm = ChatOpenAI(model="gpt-4o", temperature=0.1)

# Sample Enterprise Knowledge Base
sample_docs = [
 Document(
 page_content="Model Context Protocol (MCP) by Anthropic enables AI agents to connect to external tools over stdio and Server-Sent Events (SSE).",
 metadata={"source": "mcp_spec.md"}
 ),
 Document(
 page_content="In LangGraph, state is maintained across cyclical nodes using a StateGraph and reducer functions, allowing human-in-the-loop and retries.",
 metadata={"source": "langgraph_overview.md"}
 ),
 Document(
 page_content="Redis semantic caching stores query-vector pairs with cosine similarity thresholds to bypass redundant LLM inference calls.",
 metadata={"source": "semantic_caching.md"}
 )
]

vectorstore = Chroma.from_documents(
 documents=sample_docs,
 embedding=OpenAIEmbeddings(),
 collection_name="crag_enterprise_docs"
)
retriever = vectorstore.as_retriever(search_kwargs={"k": 3})

# Structured Grader Chains
doc_grader_prompt = ChatPromptTemplate.from_messages([
 ("system", "You are an expert retrieval evaluator. Grade whether the retrieved document is relevant to the user question. Respond with 'yes' or 'no'."),
 ("human", "Retrieved document:\n{document}\n\nUser question: {question}")
])
doc_grader = doc_grader_prompt | evaluator_llm.with_structured_output(GradeDocuments)

hallucination_grader_prompt = ChatPromptTemplate.from_messages([
 ("system", "You are a factual groundedness evaluator. Grade whether the answer is strictly supported by the provided facts. Respond with 'yes' or 'no'."),
 ("human", "Facts:\n{documents}\n\nGenerated answer: {generation}")
])
hallucination_grader = hallucination_grader_prompt | evaluator_llm.with_structured_output(GradeHallucinations)

answer_grader_prompt = ChatPromptTemplate.from_messages([
 ("system", "You are a query resolution judge. Grade whether the generated answer fully addresses the user question. Respond with 'yes' or 'no'."),
 ("human", "User question: {question}\n\nGenerated answer: {generation}")
])
answer_grader = answer_grader_prompt | evaluator_llm.with_structured_output(GradeAnswer)

from langchain_community.tools.tavily_search import TavilySearchResults

web_search_tool = TavilySearchResults(k=3)

def retrieve(state: GraphState) -> Dict[str, Any]:
 """Retrieve documents from the vector store."""
 question = state["question"]
 docs = retriever.invoke(question)
 return {"documents": docs, "question": question, "retry_count": state.get("retry_count", 0)}

def grade_documents(state: GraphState) -> Dict[str, Any]:
 """
 Evaluates retrieved chunks. Filters out low-scoring noise.
 Sets web_search flag to True if relevance drops below threshold.
 """
 question = state["question"]
 documents = state["documents"]

 filtered_docs = []
 web_search = False

 for doc in documents:
 res = doc_grader.invoke({"question": question, "document": doc.page_content})
 grade = res.binary_score
 if grade == "yes":
 filtered_docs.append(doc)
 else:
 # Irrelevant document found; trigger web search fallback to ensure completeness
 web_search = True

 # If no documents survived filtering, external search is mandatory
 if len(filtered_docs) == 0:
 web_search = True

 return {"documents": filtered_docs, "question": question, "web_search": web_search}

def transform_query(state: GraphState) -> Dict[str, Any]:
 """Transforms conversational query into an optimized web search string."""
 question = state["question"]

 re_write_prompt = ChatPromptTemplate.from_messages([
 ("system", "You are an AI query rewriter. Formulate an optimized keyword search query for Google/Tavily. Strip conversational pleasantries and return only the search string."),
 ("human", "Initial query: {question}\n\nOptimized search query:")
 ])
 rewriter = re_write_prompt | evaluator_llm
 response = rewriter.invoke({"question": question})
 better_question = response.content.strip()

 return {"documents": state["documents"], "question": better_question}

def web_search(state: GraphState) -> Dict[str, Any]:
 """Executes web search and appends external context."""
 question = state["question"]
 documents = state.get("documents", [])

 search_results = web_search_tool.invoke({"query": question})
 web_content = "\n".join([item["content"] for item in search_results])
 web_doc = Document(page_content=web_content, metadata={"source": "tavily_web_search"})

 documents.append(web_doc)
 return {"documents": documents, "question": question}

def generate(state: GraphState) -> Dict[str, Any]:
 """Generates the grounded response from filtered context."""
 question = state["question"]
 documents = state["documents"]
 retry_count = state.get("retry_count", 0)

 rag_prompt = ChatPromptTemplate.from_messages([
 ("system", "You are an expert technical assistant. Answer the question using ONLY the provided context facts. If the facts do not contain the answer, state that clearly.\n\nContext:\n{context}"),
 ("human", "Question: {question}")
 ])

 context_str = "\n\n".join([d.page_content for d in documents])
 generator_chain = rag_prompt | generator_llm
 generation = generator_chain.invoke({"context": context_str, "question": question})

 return {
 "documents": documents,
 "question": question,
 "generation": generation.content,
 "retry_count": retry_count + 1
 }

def decide_to_generate(state: GraphState) -> str:
 """Determines whether to proceed to generation or execute web search fallback."""
 web_search_flag = state["web_search"]
 if web_search_flag:
 return "transform_query"
 return "generate"

def grade_generation_v_documents_and_question(state: GraphState) -> str:
 """
 Verifies factual groundedness and query resolution before client delivery.
 Implements circuit-breaking retries (max 3).
 """
 question = state["question"]
 documents = state["documents"]
 generation = state["generation"]
 retry_count = state.get("retry_count", 1)

 if retry_count >= 3:
 return "useful" # Circuit breaker: stop recursion and deliver best effort

 context_str = "\n\n".join([d.page_content for d in documents])

 # Gate A: Groundedness check
 hallucination_result = hallucination_grader.invoke({
 "documents": context_str,
 "generation": generation
 })

 if hallucination_result.binary_score == "yes":
 # Gate B: Query relevance check
 answer_result = answer_grader.invoke({
 "question": question,
 "generation": generation
 })
 if answer_result.binary_score == "yes":
 return "useful"
 else:
 return "not_useful"
 else:
 return "not_grounded"

from langgraph.graph import END, StateGraph

workflow = StateGraph(GraphState)

# Add Nodes
workflow.add_node("retrieve", retrieve)
workflow.add_node("grade_documents", grade_documents)
workflow.add_node("generate", generate)
workflow.add_node("transform_query", transform_query)
workflow.add_node("web_search", web_search)

# Build Graph Topology
workflow.set_entry_point("retrieve")
workflow.add_edge("retrieve", "grade_documents")

workflow.add_conditional_edges(
 "grade_documents",
 decide_to_generate,
 {
 "transform_query": "transform_query",
 "generate": "generate"
 }
)

workflow.add_edge("transform_query", "web_search")
workflow.add_edge("web_search", "generate")

workflow.add_conditional_edges(
 "generate",
 grade_generation_v_documents_and_question,
 {
 "not_grounded": "generate", # Re-try generation if hallucinated
 "not_useful": "transform_query", # Reformulate search if incomplete
 "useful": END # Deliver verified answer
 }
)

# Compile Application
app = workflow.compile()

def query_crag_system(user_query: str):
 print(f"\n--- Ingesting Query: {user_query} ---")
 inputs = {"question": user_query, "retry_count": 0}

 for output in app.stream(inputs):
 for key, value in output.items():
 print(f"[State Transition] Node Finished: '{key}'")
 if "web_search" in value:
 print(f" -> Web Search Triggered: {value['web_search']}")

 print("\n--- Final Verified Response ---")
 print(value.get("generation"))

# Test 1: Internal Knowledge (Fast-Path)
query_crag_system("How does Redis semantic caching improve LLM inference costs?")

# Test 2: External Knowledge Void (Triggers Tavily Fallback)
query_crag_system("What are the key architectural changes introduced in Python 3.14 alpha release?")