agents/rag_agent.py

"""Agentic RAG implementation.

This implementation focuses on:
- Building an Agentic RAG system with dynamic search strategy
- Using LangGraph for controlling the RAG workflow
- Evaluating retrieved information quality
"""

import os.path as osp
from typing import Dict, List, Optional, Any
from langchain_core.messages import HumanMessage, AIMessage
from langchain_openai import ChatOpenAI, OpenAIEmbeddings
from langchain_core.documents import Document
from pydantic import BaseModel, Field

from langgraph.graph import StateGraph, END, START, MessagesState
from langgraph.graph.message import add_messages

# For document retrieval
from langchain_core.vectorstores import InMemoryVectorStore
from langchain_community.tools.tavily_search import TavilySearchResults
from langchain_core.tools import tool
from langchain_core.tools.retriever import create_retriever_tool
from langchain_community.document_loaders import PyPDFLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from core.chat_interface import ChatInterface
from agents.prompts import (
    DOCUMENT_EVALUATOR_PROMPT,
    DOCUMENT_SYNTHESIZER_PROMPT,
    QUERY_REWRITER_PROMPT,
)
from langchain_core.prompts import PromptTemplate
from langgraph.prebuilt import ToolNode, tools_condition
from dotenv import load_dotenv
load_dotenv()

# NOTE: Update this to the path of your documents
# For Hugging Face Spaces, use: "/home/user/app/documents/"
# For local development, use your local path
BASE_DIR = "data/"

# Add your document files here
FILE_PATHS = [
    # Example for OPM documents (you can replace with your own documents):
    osp.join(BASE_DIR, "2019-annual-performance-report.pdf"),
    osp.join(BASE_DIR, "2020-annual-performance-report.pdf"),
    osp.join(BASE_DIR, "2021-annual-performance-report.pdf"),
    osp.join(BASE_DIR, "2022-annual-performance-report.pdf"),
    

]


class DocumentEvaluation(BaseModel):
    """Evaluation result for retrieved documents."""
    is_sufficient: bool = Field(description="Whether the documents provide sufficient information")
    feedback: str = Field(description="Feedback about the document quality and what's missing")


class AgenticRAGState(MessagesState):
    """State for the Agentic RAG workflow using MessagesState as base."""
    # MessagesState already handles messages with add_messages reducer
    feedback: str = ""
    is_sufficient: bool = False
    retry_count: int = 0  # Track number of retries to prevent infinite loops
    max_retries: int = 3  # Maximum number of query rewrites allowed
    current_query_index: int = 0  # Track which message is the current query


class AgenticRAGChat(ChatInterface):
    """Agentic RAG implementation with dynamic retrieval and evaluation."""
    
    def __init__(self):
        self.llm = None
        self.embeddings = None
        self.evaluator_llm = None
        self.vector_store = None
        self.tools = []
        self.graph = None
    
    def initialize(self) -> None:
        """Initialize components for the Agentic RAG system."""
        # Initialize models
        self.llm = ChatOpenAI(model="gpt-4o", temperature=0)
        self.embeddings = OpenAIEmbeddings(model="text-embedding-3-small")
        self.evaluator_llm = self.llm.with_structured_output(DocumentEvaluation)

        # Check if documents are configured
        if FILE_PATHS and all(osp.exists(f) for f in FILE_PATHS):
            # Load documents and create vector store
            docs = self._load_and_process_documents()
            print(f"Loading {len(docs)} documents into vector store")
            self.vector_store = InMemoryVectorStore(embedding=self.embeddings)
            self.vector_store.add_documents(docs)
        else:
            print("Warning: No documents configured for RAG. Add document paths to FILE_PATHS.")
            # Create empty vector store
            self.vector_store = InMemoryVectorStore(embedding=self.embeddings)
        
        # Create tools
        self.tools = self._create_tools()
        
        # Create the graph
        self.graph = self._create_graph()

    def _load_and_process_documents(self) -> List[Document]:
        """Load and process documents for RAG."""
        docs = []
        for file_path in FILE_PATHS:
            if not osp.exists(file_path):
                print(f"Warning: File not found - {file_path}")
                continue
                
            print(f"Loading document from {file_path}")
            try:
                loader = PyPDFLoader(file_path)
                page_docs = loader.load()
                
                # Combine all pages and split into chunks
                combined_doc = "\n".join([doc.page_content for doc in page_docs])
                
                # Use RecursiveCharacterTextSplitter for better chunking
                text_splitter = RecursiveCharacterTextSplitter(
                    chunk_size=1000, 
                    chunk_overlap=200,
                    separators=["\n\n", "\n", ".", " ", ""]
                )
                chunks = text_splitter.split_text(combined_doc)
                
                # Convert chunks to Document objects with metadata
                docs.extend([
                    Document(
                        page_content=chunk, 
                        metadata={"source": osp.basename(file_path)}
                    ) for chunk in chunks
                ])
            except Exception as e:
                print(f"Error loading {file_path}: {e}")
                
        return docs
    
    def _create_tools(self) -> List[Any]:
        """Create retriever and search tools."""
        tools = []
        
        # Create retriever tool if we have documents
        if self.vector_store:
            retriever = self.vector_store.as_retriever(
                search_type="similarity",
                search_kwargs={"k": 3}
            )
            retriever_tool = create_retriever_tool(
                retriever,
                name="search_documents",
                description=(
                    "Search through the document database. "
                    "Use this for questions about content in the loaded documents."
                )
            )
            tools.append(retriever_tool)
        
        # Create web search tool
        @tool("web_search")
        def search_web(query: str) -> list[dict]:
            """
            Search the web for the latest information on any topic.
            
            Args:
                query: The search query to look up
                
            Returns:
                List of search results with title, content, and URL
            """
            search = TavilySearchResults(max_results=3)
            return search.invoke(query)
        
        tools.append(search_web)
        
        return tools

    def _generate_query_or_respond(self, state: AgenticRAGState):
        """Generate a query or respond based on the current state."""
        print("Generating query or responding...")
        
        prompt = PromptTemplate.from_template(
            """
            You are a helpful assistant that can answer questions using the provided tools.
            
            Available tools:
            - search_documents: Search through loaded documents
            - web_search: Search the web for information
            
            Based on the user's query, decide whether to use tools or respond directly.
            Use tools when you need specific information to answer the question accurately.

            Query: {question}
            """
        )
        
        # Get the latest message (either original or rewritten query)
        question = state["messages"][-1].content
        chain = prompt | self.llm.bind_tools(self.tools)
        
        response = chain.invoke({"question": question})
        return {"messages": [response]}

    def _evaluate_documents(self, state: AgenticRAGState):
        """Evaluate the documents retrieved from the retriever tool."""
        print("Evaluating documents...")
        
        # Check if we've hit max retries
        if state.get("retry_count", 0) >= state.get("max_retries", 3):
            print(f"Max retries ({state.get('max_retries', 3)}) reached. Forcing synthesis with available documents.")
            return {
                "is_sufficient": True,  # Force synthesis even if not perfect
                "feedback": "Maximum retries reached. Using available documents."
            }
        
        # Get the CURRENT user question, not the first message in history
        # Use the current_query_index to get the right message
        current_query_index = state.get("current_query_index", 0)
        
        # Find the current query message
        user_messages = [msg for msg in state["messages"] if isinstance(msg, HumanMessage)]
        if current_query_index < len(state["messages"]):
            user_question = state["messages"][current_query_index].content
        else:
            # Fallback: get the last user message
            user_question = user_messages[-1].content if user_messages else state["messages"][-1].content
        
        # Get the retrieved documents (should be the last message)
        retrieved_docs = state["messages"][-1].content
        
        print(f"Evaluating for query: '{user_question[:50]}...'")  # Debug print
        
        chain = DOCUMENT_EVALUATOR_PROMPT | self.evaluator_llm
        evaluation = chain.invoke({
            "question": user_question, 
            "retrieved_docs": retrieved_docs
        })
        
        print(f"Evaluation result: {evaluation} (Retry {state.get('retry_count', 0)}/{state.get('max_retries', 3)})")
        return {
            "is_sufficient": evaluation.is_sufficient,
            "feedback": evaluation.feedback
        }

    def _synthesize_answer(self, state: AgenticRAGState):
        """Synthesize the final answer from retrieved documents."""
        print("Synthesizing answer...")
        
        # Get the CURRENT user question using the index
        current_query_index = state.get("current_query_index", 0)
        
        # Find the current query message
        user_messages = [msg for msg in state["messages"] if isinstance(msg, HumanMessage)]
        if current_query_index < len(state["messages"]):
            user_question = state["messages"][current_query_index].content
        else:
            # Fallback: get the last user message
            user_question = user_messages[-1].content if user_messages else state["messages"][-1].content
        
        # Get the retrieved documents
        retrieved_docs = state["messages"][-1].content
        
        print(f"Synthesizing answer for: '{user_question[:50]}...'")  # Debug print
        
        chain = DOCUMENT_SYNTHESIZER_PROMPT | self.llm
        answer = chain.invoke({
            "question": user_question, 
            "retrieved_docs": retrieved_docs
        })
        
        return {"messages": [answer]}

    def _query_rewriter(self, state: AgenticRAGState):
        """Rewrite the query based on evaluation feedback."""
        print("Rewriting query...")
        
        # Increment retry count
        current_retry = state.get("retry_count", 0)
        
        # Get the CURRENT user question using the index
        current_query_index = state.get("current_query_index", 0)
        
        # Find the current query message
        user_messages = [msg for msg in state["messages"] if isinstance(msg, HumanMessage)]
        if current_query_index < len(state["messages"]):
            user_question = state["messages"][current_query_index].content
        else:
            # Fallback: get the last user message
            user_question = user_messages[-1].content if user_messages else state["messages"][-1].content
        
        retrieved_docs = state["messages"][-1].content
        feedback = state["feedback"]
        
        print(f"Rewriting query for: '{user_question[:50]}...'")  # Debug print
        
        chain = QUERY_REWRITER_PROMPT | self.llm
        new_query = chain.invoke({
            "question": user_question, 
            "feedback": feedback,
            "retrieved_docs": retrieved_docs
        })
        
        print(f"Rewritten query (Attempt {current_retry + 1}/{state.get('max_retries', 3)}): {new_query.content}")
        return {
            "messages": [new_query],
            "retry_count": current_retry + 1  # Increment retry count
        }
    
    def _create_graph(self) -> Any:
        """Create the agentic RAG graph."""
        # Create the graph builder
        graph_builder = StateGraph(AgenticRAGState)
        
        # Add nodes
        graph_builder.add_node("generate_query_or_respond", self._generate_query_or_respond)
        graph_builder.add_node("retrieve_documents", ToolNode(self.tools))
        graph_builder.add_node("evaluate_documents", self._evaluate_documents)
        graph_builder.add_node("synthesize_answer", self._synthesize_answer)
        graph_builder.add_node("query_rewriter", self._query_rewriter)

        # Add edges
        graph_builder.add_edge(START, "generate_query_or_respond")
        
        # Conditional edge: if tools were called, retrieve documents; else end
        graph_builder.add_conditional_edges(
            "generate_query_or_respond",
            tools_condition,
            {
                "tools": "retrieve_documents",
                END: END,
            },
        )
        
        # After retrieval, evaluate documents
        graph_builder.add_edge("retrieve_documents", "evaluate_documents")
        
        # Conditional edge: if sufficient, synthesize; else rewrite query
        graph_builder.add_conditional_edges(
            "evaluate_documents",
            lambda x: "synthesize_answer" if x["is_sufficient"] else "query_rewriter",
            {
                "synthesize_answer": "synthesize_answer",
                "query_rewriter": "query_rewriter",
            },
        )
        
        # After rewriting, generate new query
        graph_builder.add_edge("query_rewriter", "generate_query_or_respond")
        
        # After synthesizing, end
        graph_builder.add_edge("synthesize_answer", END)
        
        return graph_builder.compile()
    
    
    def _convert_history_to_messages(self, chat_history: Optional[List[Dict[str, str]]]) -> List:
        """Convert chat history to LangChain message format.
        
        Args:
            chat_history: List of dicts with 'role' and 'content' keys
            
        Returns:
            List of LangChain message objects
        """
        messages = []
        if chat_history:
            for msg in chat_history:
                if msg["role"] == "user":
                    messages.append(HumanMessage(content=msg["content"]))
                elif msg["role"] == "assistant":
                    messages.append(AIMessage(content=msg["content"]))
        return messages
    
    
    
    def process_message(self, message: str, chat_history: Optional[List[Dict[str, str]]] = None) -> str:
        """Process a message using the Agentic RAG system."""
        print("\n=== STARTING AGENTIC RAG QUERY ===")
        print(f"Query: {message}")
        
        # Convert chat history to messages
        history_messages = self._convert_history_to_messages(chat_history)
    
        # Mark the position where the current query starts
        # This is important for the evaluator to know which is the actual query
        history_length = len(history_messages)
        
        # Add the current message
        current_query_message = HumanMessage(content=message)
        history_messages.append(current_query_message)
        
        # Create initial state with full conversation history
        # Store the index of the current query for reference
        state = AgenticRAGState(
            messages=history_messages,
            feedback="",
            is_sufficient=False,
            retry_count=0,
            max_retries=3,
            # Add this to track the current query index
            current_query_index=history_length  # This is the index of the current query
        )
        
        try:
            # Run the workflow with increased recursion limit
            config = {"recursion_limit": 30}
            result = self.graph.invoke(state, config=config)
            
            print("\n=== RAG QUERY COMPLETED ===")
            
            # Return the final answer
            if result.get("messages"):
                final_message = result["messages"][-1]
                if hasattr(final_message, 'content'):
                    return final_message.content
                else:
                    return str(final_message)
            else:
                return "I couldn't find relevant information to answer your question."
                
        except Exception as e:
            print(f"Error in RAG processing: {e}")
            if "recursion" in str(e).lower():
                return ("I had difficulty finding the exact information you're looking for in the documents. "
                    "Based on the available documents, I can see references to various topics, "
                    "but I couldn't find specific details. You might want to try asking about a specific aspect.")
            return f"I encountered an error while searching for information: {str(e)}"