Stateful Agentic RAG: Graph State Machines, Self-Correction Loops & Fallback Routing

# rag_graph/state.py
from typing import List, Optional, Literal
from langchain_core.documents import Document
from langchain_core.messages import BaseMessage, HumanMessage, AIMessage
from langgraph.graph import StateGraph, END
from pydantic import BaseModel, Field

class AgentState(BaseModel):
    """
    Represents the state of our RAG agent.
    This state is passed between nodes in the graph.
    """
    query: str = Field(description="The original user query.")
    chat_history: List[BaseMessage] = Field(default_factory=list, description="Full chat history.")
    documents: List[Document] = Field(default_factory=list, description="Retrieved documents.")
    generation: Optional[str] = Field(None, description="Generated LLM response.")
    retrieval_attempts: int = Field(0, description="Number of retrieval attempts.")
    web_search_performed: bool = Field(False, description="Flag indicating if web search was performed.")
    # Add a field to track the current decision path for debugging/logging
    current_path: List[str] = Field(default_factory=list, description="Path taken through the graph.")

    class Config:
        arbitrary_types_allowed = True # Allow BaseMessage

# rag_graph/nodes.py
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnablePassthrough
from langchain_openai import ChatOpenAI
from langchain_community.vectorstores import FAISS
from langchain_openai import OpenAIEmbeddings
from langchain_community.tools import DuckDuckGoSearchRun
from langchain_core.output_parsers import StrOutputParser
from typing import List, Dict, Any

# Assume these are initialized globally or passed in
# For production, use environment variables for API keys and proper vector store setup
llm = ChatOpenAI(model="gpt-4o", temperature=0)
embeddings = OpenAIEmbeddings()
vectorstore = FAISS.from_texts(["LangGraph is a library for building stateful, multi-actor applications with LLMs.",
                                "RAG stands for Retrieval Augmented Generation.",
                                "Self-correction in RAG improves accuracy.",
                                "Postgres can be used for state persistence."], embeddings)
retriever = vectorstore.as_retriever()
web_search_tool = DuckDuckGoSearchRun()

# --- Prompts ---
ROUTER_PROMPT = ChatPromptTemplate.from_messages([
    ("system", "You are a smart routing agent. Your goal is to determine the best next step for a user query."),
    ("human", """Given the user query and chat history, decide whether to:
    1. 'retrieve': Perform a standard RAG retrieval from our internal knowledge base.
    2. 'web_search': Perform a web search for external information.
    3. 'generate': Directly generate a response if the query is simple and doesn't require retrieval or search.

    Respond with only one of the keywords: 'retrieve', 'web_search', 'generate'.

    Chat History: {chat_history}
    User Query: {query}
    """)
])

RETRIEVAL_GRADER_PROMPT = ChatPromptTemplate.from_messages([
    ("system", "You are a document relevance grader. Your task is to assess if the retrieved documents are relevant to the user's query."),
    ("human", """User Query: {query}
    Retrieved Documents: {documents}

    Are the retrieved documents relevant to the user's query?
    Respond with 'yes' or 'no'.
    """)
])

HALLUCINATION_GRADER_PROMPT = ChatPromptTemplate.from_messages([
    ("system", "You are a hallucination grader. Your task is to assess if the generated answer is grounded in the provided documents."),
    ("human", """User Query: {query}
    Retrieved Documents: {documents}
    Generated Answer: {generation}

    Is the generated answer fully supported by the retrieved documents?
    Respond with 'yes' or 'no'.
    """)
])

ANSWER_GENERATION_PROMPT = ChatPromptTemplate.from_messages([
    ("system", "You are an AI assistant. Use the following retrieved context to answer the user's question. If the context does not contain the answer, state that you don't know."),
    ("human", """Context: {documents}
    Chat History: {chat_history}
    Question: {query}
    """)
])

# --- Nodes ---
def route_query(state: AgentState) -> str:
    """Decides the next step based on the query."""
    print("---ROUTE QUERY---")
    state.current_path.append("route_query")
    router_chain = ROUTER_PROMPT | llm | StrOutputParser()
    decision = router_chain.invoke({"query": state.query, "chat_history": state.chat_history})
    print(f"Router decision: {decision}")
    if "web_search" in decision.lower():
        return "web_search"
    elif "retrieve" in decision.lower():
        return "retrieve"
    else: # Default to generate if not explicitly retrieve or web_search
        return "generate"

def retrieve(state: AgentState) -> AgentState:
    """Retrieves documents from the vector store."""
    print("---RETRIEVE DOCUMENTS---")
    state.current_path.append("retrieve")
    state.retrieval_attempts += 1
    documents = retriever.invoke(state.query)
    state.documents = documents
    return state

def grade_documents(state: AgentState) -> str:
    """Grades the relevance of retrieved documents."""
    print("---GRADE DOCUMENTS---")
    state.current_path.append("grade_documents")
    if not state.documents:
        print("No documents retrieved, initiating web search.")
        return "no_documents"

    grader_chain = RETRIEVAL_GRADER_PROMPT | llm | StrOutputParser()
    decision = grader_chain.invoke({"query": state.query, "documents": state.documents})
    print(f"Document relevance decision: {decision}")
    if "yes" in decision.lower():
        print("Documents are relevant.")
        return "relevant"
    else:
        print("Documents are not relevant, initiating web search.")
        return "not_relevant"

def web_search(state: AgentState) -> AgentState:
    """Performs a web search and adds results to documents."""
    print("---WEB SEARCH---")
    state.current_path.append("web_search")
    state.web_search_performed = True
    web_results = web_search_tool.invoke({"query": state.query})
    # Convert web results to Document objects
    state.documents.extend([Document(page_content=web_results, metadata={"source": "web_search"})])
    return state

def generate_response(state: AgentState) -> AgentState:
    """Generates a final answer based on retrieved documents and query."""
    print("---GENERATE RESPONSE---")
    state.current_path.append("generate_response")
    # Format documents for the prompt
    docs_content = "\n\n".join([doc.page_content for doc in state.documents])

    generation_chain = ANSWER_GENERATION_PROMPT | llm | StrOutputParser()
    response = generation_chain.invoke({
        "query": state.query,
        "documents": docs_content,
        "chat_history": state.chat_history
    })
    state.generation = response
    return state

def grade_generation_for_hallucination(state: AgentState) -> str:
    """Grades the generated answer for hallucination against retrieved documents."""
    print("---GRADE GENERATION FOR HALLUCINATION---")
    state.current_path.append("grade_generation_for_hallucination")
    if not state.documents: # If no documents were used, we can't grade against them
        print("No documents to grade against, assuming no hallucination for now.")
        return "no_hallucination"

    grader_chain = HALLUCINATION_GRADER_PROMPT | llm | StrOutputParser()
    decision = grader_chain.invoke({
        "query": state.query,
        "documents": state.documents,
        "generation": state.generation
    })
    print(f"Hallucination decision: {decision}")
    if "yes" in decision.lower():
        print("Generation is grounded in documents.")
        return "no_hallucination"
    else:
        print("Generation contains hallucination, attempting web search fallback.")
        return "hallucination"

def update_chat_history(state: AgentState) -> AgentState:
    """Updates the chat history with the latest query and response."""
    print("---UPDATE CHAT HISTORY---")
    state.current_path.append("update_chat_history")
    state.chat_history.append(HumanMessage(content=state.query))
    if state.generation:
        state.chat_history.append(AIMessage(content=state.generation))
    return state

# rag_graph/graph.py
from langgraph.checkpoint.sqlite import SqliteSaver # For local testing
from langgraph.checkpoint.postgres import PostgresSaver # For production
from langgraph.graph import StateGraph, END
from rag_graph.state import AgentState
from rag_graph.nodes import (
    route_query, retrieve, grade_documents, web_search,
    generate_response, grade_generation_for_hallucination, update_chat_history
)
import os

# For production, configure PostgresSaver
# memory = PostgresSaver.from_conn_string(os.environ["POSTGRES_CONNECTION_STRING"])
# For local testing, use SqliteSaver
memory = SqliteSaver.from_conn_string(":memory:") # In-memory SQLite for quick testing

def build_graph():
    workflow = StateGraph(AgentState)

    # Define nodes
    workflow.add_node("retrieve", retrieve)
    workflow.add_node("grade_documents", grade_documents)
    workflow.add_node("web_search", web_search)
    workflow.add_node("generate_response", generate_response)
    workflow.add_node("grade_generation_for_hallucination", grade_generation_for_hallucination)
    workflow.add_node("update_chat_history", update_chat_history)

    # Set entry point
    workflow.set_entry_point("route_query")

    # Define edges
    workflow.add_conditional_edges(
        "route_query",
        route_query,
        {
            "retrieve": "retrieve",
            "web_search": "web_search",
            "generate": "generate_response" # Direct generation for simple queries
        }
    )

    workflow.add_edge("retrieve", "grade_documents")

    workflow.add_conditional_edges(
        "grade_documents",
        grade_documents,
        {
            "relevant": "generate_response",
            "not_relevant": "web_search",
            "no_documents": "web_search" # If retrieval yielded nothing, try web search
        }
    )

    workflow.add_edge("web_search", "generate_response") # After web search, always try to generate

    workflow.add_edge("generate_response", "grade_generation_for_hallucination")

    workflow.add_conditional_edges(
        "grade_generation_for_hallucination",
        grade_generation_for_hallucination,
        {
            "no_hallucination": "update_chat_history",
            "hallucination": "web_search" # If hallucination, try web search (if not already done)
        }
    )

    workflow.add_edge("update_chat_history", END)

    # Compile the graph
    app = workflow.compile(checkpointer=memory)
    return app

# Example usage (in a separate script or main block)
if __name__ == "__main__":
    app = build_graph()

    # Example 1: Simple RAG query
    print("\n--- Running Example 1: Simple RAG ---")
    config = {"configurable": {"thread_id": "1"}}
    inputs = {"query": "What is LangGraph?", "chat_history": []}
    for s in app.stream(inputs, config=config):
        print(s)
    final_state = app.get_state(config)
    print(f"\nFinal Answer (Thread 1): {final_state.values['generation']}")
    print(f"Path taken (Thread 1): {final_state.values['current_path']}")

    # Example 2: Query requiring web search (e.g., current events)
    print("\n--- Running Example 2: Web Search Fallback ---")
    config = {"configurable": {"thread_id": "2"}}
    inputs = {"query": "What is the capital of France?", "chat_history": []} # Assume internal KB doesn't have this
    for s in app.stream(inputs, config=config):
        print(s)
    final_state = app.get_state(config)
    print(f"\nFinal Answer (Thread 2): {final_state.values['generation']}")
    print(f"Path taken (Thread 2): {final_state.values['current_path']}")

    # Example 3: Query that might lead to hallucination or irrelevant docs
    print("\n--- Running Example 3: Hallucination/Irrelevant Docs ---")
    config = {"configurable": {"thread_id": "3"}}
    inputs = {"query": "Tell me about the latest advancements in quantum computing, specifically related to cold fusion.", "chat_history": []}
    for s in app.stream(inputs, config=config):
        print(s)
    final_state = app.get_state(config)
    print(f"\nFinal Answer (Thread 3): {final_state.values['generation']}")
    print(f"Path taken (Thread 3): {final_state.values['current_path']}")