Self_Corrective_RAG_Agent.py

# ## Importing dependencies
import re
from typing import Literal, TypedDict
import openai
from timescale_vector import client  # Client for storing and retrieving vector embeddings from the Timescale/Postgres DB
from pydantic import BaseModel, Field
import os
from dotenv import load_dotenv

from langchain import hub

from langchain_community.tools.tavily_search import TavilySearchResults

from langchain_core.documents import Document
from langchain_core.output_parsers import StrOutputParser
from langchain_core.messages import BaseMessage, AIMessage, convert_to_messages
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.retrievers import BaseRetriever

from langchain_openai import ChatOpenAI
from langchain_openai import OpenAIEmbeddings

from langgraph.graph import END, StateGraph, MessagesState



load_dotenv()  # Load environment variables from a .env file

# Access API keys and credentials
OPENAI_API_KEY    = os.environ["OPENAI_API_KEY"]
TIMESCALE_DB_URI  = os.environ["TIMESCALE_DB_URI"]
TAVILY_API_KEY    = os.environ["TAVILY_API_KEY"]
LANGCHAIN_API_KEY = os.environ["LANGCHAIN_API_KEY"] 


# Example usage of OpenAI embedding (commented out for reference):
# query = "Explain The uncertainty principle in Quantum Mechanics"
# query_embedding = openai.embeddings.create(
#     input=[query],  # Input can be a string or list of strings
#     model="text-embedding-3-small"
# )

# Configuration for Timescale vector storage
table_name     = 'embeddings_table'  # Table name for storing vector embeddings
embedding_dims = 1536                # Embedding dimensionality for 'text-embedding-3-small'

# Initialize vector client to manage vector store operations
vec_client = client.Sync(
    TIMESCALE_DB_URI,
    table_name,
    embedding_dims
)


# Regular expression to clean newline characters
NEWLINE_RE = re.compile("\n+")

def get_docs(question: str, documents: list[Document]) -> list[Document]:
    """
    Retrieve top-k similar documents based on vector similarity for a given question.

    Args:
        question (str): The user's question or query.
        documents (list[Document]): Existing list of documents (can be empty). 
                                    New retrieved documents will be appended here.

    Returns:
        list[Document]: The updated list with newly retrieved documents appended.
    """
    # Generate embedding for the user's query
    query_embedding = openai.embeddings.create(
        input=[question],  # Accepts a string or list of strings
        model="text-embedding-3-small"
    )

    # Retrieve top 5 most similar embeddings from vector store
    results = vec_client.search(query_embedding.data[0].embedding, limit=5)

    # Convert DB rows into LangChain Document objects and append to list
    for row in results:
        documents.append(Document(page_content=row[2], metadata=row[1]))
    
    return documents


# Initialize the main language model
llm = ChatOpenAI(model="gpt-4.1-nano-2025-04-14", temperature=0)

# Initialize Tavily web search tool
tavily_search_tool = TavilySearchResults(max_results=3)


# Define the graph state shared across nodes
class GraphState(MessagesState):
    question:         str              # User's question
    documents:        list[Document]   # Retrieved documents relevant to the question
    candidate_answer: str              # Answer generated by the LLM
    retries:          int              # Retry count for hallucination/irrelevant answers
    web_fallback:     bool             # Whether to fall back to web search if retries are exhausted

class GraphConfig(TypedDict):
    max_retries: int  # Maximum number of retries before using web search


MAX_RETRIES = 3
VERBOSE     = True


# ----------------------
# Graph Node: Document Search
# ----------------------
def document_search(state: GraphState, human_message: int = -3):
    """
    Retrieve relevant documents from the vector database based on the user's question.

    Args:
        state (dict): Current graph state.

    Returns:
        dict: Updated state with the retrieved documents and original question.
    """
    if VERBOSE:
        print("---RETRIEVE---")

    # Extract question from the message history
    # if this is a subgraph the human question will be 3rd last in the messages list
    try:
        question = str(convert_to_messages(state["messages"])[human_message].content)
    # if this is the main graph the human question will be the last one in the messages list
    except:
        question = str(convert_to_messages(state["messages"])[-1].content)
        
    # Retrieve documents using embedding similarity
    documents = get_docs(question, state.get("documents", []))
    

    return {"documents": documents, "question": question, "web_fallback": True}


# ----------------------
# Graph Node: Answer Generation
# ----------------------
RAG_PROMPT: ChatPromptTemplate = hub.pull("rlm/rag-prompt")
def generate(state: GraphState):
    """
    Generate an answer using the retrieved documents and user's question.

    Args:
        state (dict): Current graph state.

    Returns:
        dict: Updated state with a candidate answer and incremented retry count.
    """
    if VERBOSE:
        print("---GENERATE---")
    
    question  = state["question"]
    documents = state["documents"]
    retries   = state["retries"] if state.get("retries") is not None else -1

    rag_chain  = RAG_PROMPT | llm | StrOutputParser()
    generation = rag_chain.invoke({"context": documents, "question": question})
    
    return {"retries": retries + 1, "candidate_answer": generation}


# ----------------------
# Graph Node: Query Rewriting
# ----------------------
QUERY_REWRITER_SYSTEM = (
"""
You are a question rewriter that converts an input question into an optimized version 
for vector store retrieval. Understand the semantic intent of the question.
"""
)

QUERY_REWRITER_PROMPT = ChatPromptTemplate.from_messages(
    [
        ("system", QUERY_REWRITER_SYSTEM),
        (
            "human",
            "Here is the initial question: \n\n {question} \n Formulate an improved question.",
        ),
    ]
)

def transform_query(state: GraphState):
    """
    Rewrite the question to improve retrieval performance.

    Args:
        state (dict): Current graph state.

    Returns:
        dict: Updated state with a reformulated question.
    """
    if VERBOSE:
        print("---TRANSFORM QUERY---")
    
    question = state["question"]
    query_rewriter  = QUERY_REWRITER_PROMPT | llm | StrOutputParser()
    better_question = query_rewriter.invoke({"question": question})

    return {"question": better_question}


# ----------------------
# Graph Node: Web Search Fallback
# ----------------------
def web_search(state: GraphState):
    """
    Perform a web search as a fallback mechanism when local document retrieval fails.

    Args:
        state (dict): Current graph state.

    Returns:
        dict: Updated state with documents retrieved from web search.
    """
    if VERBOSE:
        print("---RUNNING WEB SEARCH---")

    question  = state["question"]
    documents = state["documents"]
    
    search_results = tavily_search_tool.invoke(question)
    search_content = "\n".join([d["content"] for d in search_results])

    documents.append(Document(page_content=search_content, metadata={"source": "websearch"}))
    
    return {"documents": documents, "web_fallback": False}


# ----------------------
# Graph Node: Finalize Answer
# ----------------------
def finalize_response(state: GraphState):
    """
    Finalize and return the generated response to the user.

    Args:
        state (dict): Current graph state.

    Returns:
        dict: Updated state with final AI message.
    """
    if VERBOSE:
        print("---FINALIZING THE RESPONSE---")
    
    return {"messages": [AIMessage(content=state["candidate_answer"])]}


# ----------------------
# Answer Grading - Hallucination & Relevance
# ----------------------

class GradeHallucinations(BaseModel):
    """Determine whether the generated answer is grounded in provided documents."""
    binary_score: str = Field(description="Answer is grounded in the facts, 'yes' or 'no'")


HALLUCINATION_GRADER_SYSTEM = (
"""
You are a grader assessing whether an LLM's generation is grounded in a set of retrieved facts.
Return 'yes' if fully supported by the facts; otherwise, return 'no'.

If the generation includes code, ensure it exactly matches the retrieved facts.
"""
)

HALLUCINATION_GRADER_PROMPT = ChatPromptTemplate.from_messages([
    ("system", HALLUCINATION_GRADER_SYSTEM),
    ("human", "Set of facts: \n\n {documents} \n\n LLM generation: {generation}"),
])


class GradeAnswer(BaseModel):
    """Determine whether the generated answer addresses the user's question."""
    binary_score: str = Field(description="Answer addresses the question, 'yes' or 'no'")


ANSWER_GRADER_SYSTEM = (
"""
You are a grader determining if an answer resolves a user's question.
Return 'yes' if it does, otherwise return 'no'.
"""
)

ANSWER_GRADER_PROMPT = ChatPromptTemplate.from_messages([
    ("system", ANSWER_GRADER_SYSTEM),
    ("human", "User question: \n\n {question} \n\n LLM generation: {generation}"),
])


def grade_generation_v_documents_and_question(state: GraphState, config) -> Literal["generate", "transform_query", "web_search", "finalize_response"]:
    """
    Decide the next step based on whether the generated answer is both grounded in documents and relevant to the question.

    Args:
        state (dict): Current graph state.
        config (dict): Graph configuration containing retry logic.

    Returns:
        str: Next node to call in the graph.
    """
    question     = state["question"]
    documents    = state["documents"]
    generation   = state["candidate_answer"]
    web_fallback = state["web_fallback"]
    retries      = state["retries"] if state.get("retries") is not None else -1
    max_retries  = config.get("configurable", {}).get("max_retries", MAX_RETRIES)

    if not web_fallback:
        return "finalize_response"

    if VERBOSE:
        print("---CHECK HALLUCINATIONS---")

    hallucination_grader = HALLUCINATION_GRADER_PROMPT | llm.with_structured_output(GradeHallucinations)
    hallucination_grade: GradeHallucinations = hallucination_grader.invoke(
        {"documents": documents, "generation": generation}
    )

    if hallucination_grade.binary_score == "no":
        if VERBOSE: print("---DECISION: ANSWER NOT GROUNDED, RETRY---")
        return "generate" if retries < max_retries else "web_search"

    if VERBOSE:
        print("---DECISION: ANSWER GROUNDED, NOW CHECK RELEVANCE---")

    answer_grader = ANSWER_GRADER_PROMPT | llm.with_structured_output(GradeAnswer)
    answer_grade: GradeAnswer = answer_grader.invoke({"question": question, "generation": generation})
    
    if answer_grade.binary_score == "yes":
        if VERBOSE: print("---DECISION: ANSWER RELEVANT---")
        return "finalize_response"
    else:
        if VERBOSE: print("---DECISION: ANSWER IRRELEVANT, RETRY---")
        return "transform_query" if retries < max_retries else "web_search"


# ----------------------
# Graph Construction
# ----------------------
self_corrective_rag_builder = StateGraph(GraphState, config_schema=GraphConfig)

# Add graph nodes
self_corrective_rag_builder.add_node("document_search",   document_search)
self_corrective_rag_builder.add_node("generate",          generate)
self_corrective_rag_builder.add_node("transform_query",   transform_query)
self_corrective_rag_builder.add_node("web_search",        web_search)
self_corrective_rag_builder.add_node("finalize_response", finalize_response)

# Define graph edges
self_corrective_rag_builder.set_entry_point("document_search")
self_corrective_rag_builder.add_edge("document_search",   "generate")
self_corrective_rag_builder.add_edge("transform_query",   "document_search")
self_corrective_rag_builder.add_edge("web_search",        "generate")
self_corrective_rag_builder.add_edge("finalize_response", END)
self_corrective_rag_builder.add_conditional_edges("generate", grade_generation_v_documents_and_question)


self_corrective_rag_agent = self_corrective_rag_builder.compile()

# # Compile the graph
# self_crag = self_corrective_rag_builder.compile()


# # ----------------------
# # Visualize the graph
# # ----------------------
# from IPython.display import Image, display
# display(Image(self_crag.get_graph().draw_mermaid_png()))


# Example usage:
# VERBOSE = True
# inputs = {"messages": [("human", "explain uncertainty principle in quantum mechanics")]}
# for output in self_crag.stream(inputs):
#     print("\n---\n")

# VERBOSE = False
# inputs = {"messages": [("human", "explain diffraction in crystals")]}
# for output in self_crag.stream(inputs):
#     print(output)
#     print("\n---\n")