app/llm/graph.py

"""
LangGraph-based multi-agent workflow for Builder's AI.
This implements a graph-based orchestration of multiple specialized agents.
"""
from typing import Dict, List, Optional
from langgraph.graph import StateGraph, END
import json

from app.llm.state import AgentState
from app.llm.agents.router import router_agent
from app.llm.agents.search import search_agent
from app.llm.agents.rag import rag_agent
from app.llm.agents.policy import policy_agent
from app.llm.agents.general import general_agent
from app.services.rag_service import rag_service
from app.utils.embeddings import embedding_generator


class MultiAgentGraph:
    """LangGraph-based multi-agent workflow orchestrator."""
    
    def __init__(self):
        """Initialize the multi-agent graph."""
        self.graph = self._build_graph()
        print("[Multi-Agent Graph] Initialized")
    
    def _build_graph(self) -> StateGraph:
        """
        Build the LangGraph workflow.
        
        Returns:
            Compiled StateGraph
        """
        # Create workflow graph
        workflow = StateGraph(AgentState)
        
        # Add nodes
        workflow.add_node("router", self._router_node)
        workflow.add_node("search_agent", self._search_node)
        workflow.add_node("rag_agent", self._rag_node)
        workflow.add_node("policy_agent", self._policy_node)
        workflow.add_node("general_agent", self._general_node)
        
        # Set entry point
        workflow.set_entry_point("router")
        
        # Add conditional edges from router to specialized agents
        workflow.add_conditional_edges(
            "router",
            self._route_query,
            {
                "search": "search_agent",
                "rag": "rag_agent",
                "policy": "policy_agent",
                "general": "general_agent"
            }
        )
        
        # All agent nodes end the workflow
        workflow.add_edge("search_agent", END)
        workflow.add_edge("rag_agent", END)
        workflow.add_edge("policy_agent", END)
        workflow.add_edge("general_agent", END)
        
        # Compile the graph
        return workflow.compile()
    
    def _router_node(self, state: AgentState) -> AgentState:
        """
        Router node: Determines which specialized agent should handle the query.
        
        Args:
            state: Current agent state
            
        Returns:
            Updated state with routing decision
        """
        print(f"[Router Node] Processing query: {state['query'][:50]}...")
        
        try:
            # Use router agent to determine the appropriate agent
            routing = router_agent.route(
                query=state["query"],
                chat_history=state.get("chat_history", [])
            )
            
            agent_type = routing.get("agent", "general")
            reasoning = routing.get("reasoning", "")
            
            print(f"[Router Node] Routing to: {agent_type} - {reasoning}")
            
            return {
                **state,
                "agent_type": agent_type,
                "routing_reasoning": reasoning
            }
            
        except Exception as e:
            print(f"[Router Node] Error: {e}")
            return {
                **state,
                "agent_type": "general",
                "routing_reasoning": f"Error in routing: {str(e)}",
                "error": str(e)
            }
    
    def _route_query(self, state: AgentState) -> str:
        """
        Conditional edge function to route to the appropriate agent.
        
        Args:
            state: Current agent state
            
        Returns:
            Agent type string
        """
        return state.get("agent_type", "general")
    
    def _search_node(self, state: AgentState) -> AgentState:
        """
        Search agent node: Performs web search and generates answer.
        
        Args:
            state: Current agent state
            
        Returns:
            Updated state with search results and answer
        """
        print("[Search Node] Executing web search...")
        
        try:
            response = search_agent.search_and_answer(state["query"])
            
            return {
                **state,
                "answer": response.get("answer", ""),
                "sources": response.get("sources", []),
                "search_results": response.get("sources", []),
                "metadata": {
                    "agent": "search",
                    "routing_reasoning": state.get("routing_reasoning", "")
                }
            }
            
        except Exception as e:
            print(f"[Search Node] Error: {e}")
            return {
                **state,
                "answer": "I encountered an error while searching. Please try again.",
                "sources": [],
                "error": str(e)
            }
    
    def _rag_node(self, state: AgentState) -> AgentState:
        """
        RAG agent node: Retrieves relevant documents and generates answer.
        
        Args:
            state: Current agent state
            
        Returns:
            Updated state with RAG context and answer
        """
        print("[RAG Node] Performing semantic search...")
        
        try:
            # Check if policy IDs are provided
            policy_ids = state.get("policy_ids")
            
            if policy_ids:
                # Search within selected policies
                print(f"[RAG Node] Searching within {len(policy_ids)} selected policies")
                print(f"[RAG Node] Policy IDs: {policy_ids}")
                
                context_chunks = rag_service.search_policies(
                    query=state["query"],
                    policy_ids=policy_ids,
                    top_k=10  # Increased for better coverage
                )
                
                print(f"[RAG Node] Found {len(context_chunks)} chunks from policies")
            else:
                # Regular document search
                print(f"[RAG Node] Searching user documents for user_id: {state.get('user_id')}")
                context_chunks = rag_service.semantic_search(
                    query=state["query"],
                    user_id=state.get("user_id"),
                    top_k=10  # Increased for better coverage
                )
                print(f"[RAG Node] Found {len(context_chunks)} chunks from user docs")
            
            if not context_chunks:
                print("[RAG Node] No relevant documents found")
                no_doc_message = (
                    "I don't have any content in the selected policies to answer this question."
                    if policy_ids
                    else "I don't have any uploaded documents to answer this question. Please upload construction documents or ask a general question."
                )
                return {
                    **state,
                    "answer": no_doc_message,
                    "sources": [],
                    "context_chunks": [],
                    "metadata": {
                        "agent": "rag",
                        "note": "No documents available",
                        "policy_mode": bool(policy_ids)
                    }
                }
            
            # Generate answer using RAG agent
            response = rag_agent.answer(state["query"], context_chunks)
            
            # Determine agent label: "policy" if searching official policies, "rag" if user docs
            agent_label = "policy" if policy_ids else "rag"
            
            return {
                **state,
                "answer": response.get("answer", ""),
                "sources": response.get("sources", []),
                "context_chunks": context_chunks,
                "metadata": {
                    "agent": agent_label,  # "policy" or "rag"
                    "chunks_retrieved": len(context_chunks),
                    "routing_reasoning": state.get("routing_reasoning", ""),
                    "policy_mode": bool(policy_ids),
                    "policy_count": len(policy_ids) if policy_ids else 0
                }
            }
            
        except Exception as e:
            print(f"[RAG Node] Error: {e}")
            return {
                **state,
                "answer": "I encountered an error while processing your document query. Please try again.",
                "sources": [],
                "error": str(e)
            }
    
    def _policy_node(self, state: AgentState) -> AgentState:
        """
        Policy agent node: Handles regulatory and compliance queries using official policy documents.
        
        Args:
            state: Current agent state
            
        Returns:
            Updated state with policy answer
        """
        print("[Policy Node] Processing policy query...")
        
        try:
            # Check if policies are selected
            policy_ids = state.get("policy_ids", [])
            
            if not policy_ids:
                print("[Policy Node] No policies selected, redirecting to RAG agent")
                return {
                    **state,
                    "answer": "Please select at least one policy document from the sidebar to get policy-specific answers.",
                    "sources": [],
                    "metadata": {
                        "agent": "policy",
                        "note": "No policies selected",
                        "routing_reasoning": state.get("routing_reasoning", "")
                    }
                }
            
            # Search selected official policies for relevant information
            policy_filter = {
                "$and": [
                    {"user_id": {"$eq": "official_policies"}},
                    {"document_id": {"$in": policy_ids}}
                ]
            }
            
            context_chunks = rag_service.collection.query(
                query_embeddings=[embedding_generator.generate_embedding(state["query"])],
                n_results=10,
                where=policy_filter
            )
            
            # Format chunks
            if context_chunks and context_chunks['documents']:
                formatted_chunks = [
                    {
                        "content": context_chunks['documents'][0][i],
                        "metadata": context_chunks['metadatas'][0][i]
                    }
                    for i in range(len(context_chunks['documents'][0]))
                ]
            else:
                formatted_chunks = []
            
            if not formatted_chunks:
                return {
                    **state,
                    "answer": "I couldn't find relevant information in the selected policy documents. Please try rephrasing your question or selecting different policies.",
                    "sources": [],
                    "metadata": {
                        "agent": "policy",
                        "note": "No relevant content found in selected policies"
                    }
                }
            
            # Use policy agent with context
            response = policy_agent.answer(state["query"], formatted_chunks)
            
            print(f"[Policy Node] Response policy_names: {response.get('policy_names', [])}")
            
            return {
                **state,
                "answer": response.get("answer", ""),
                "sources": response.get("sources", []),
                "policy_names": response.get("policy_names", []),  # Pass policy names through
                "metadata": {
                    "agent": "policy",
                    "routing_reasoning": state.get("routing_reasoning", ""),
                    "chunks_retrieved": len(formatted_chunks),
                    "policy_names": response.get("policy_names", [])  # Include in metadata too
                }
            }
            
        except Exception as e:
            print(f"[Policy Node] Error: {e}")
            return {
                **state,
                "answer": "I encountered an error while processing your policy question. Please try again.",
                "sources": [],
                "error": str(e)
            }
    
    def _general_node(self, state: AgentState) -> AgentState:
        """
        General agent node: Handles general construction questions and conversations.
        
        Args:
            state: Current agent state
            
        Returns:
            Updated state with general answer
        """
        print("[General Node] Processing general query...")
        
        try:
            # Format chat history for the agent
            chat_history = state.get("chat_history", [])
            
            response = general_agent.answer(
                query=state["query"],
                chat_history=chat_history
            )
            
            return {
                **state,
                "answer": response.get("answer", ""),
                "sources": [],
                "metadata": {
                    "agent": "general",
                    "routing_reasoning": state.get("routing_reasoning", "")
                }
            }
            
        except Exception as e:
            print(f"[General Node] Error: {e}")
            return {
                **state,
                "answer": "I apologize, but I encountered an error. Please try again.",
                "sources": [],
                "error": str(e)
            }
    
    def process_query(
        self,
        query: str,
        user_id: Optional[str] = None,
        chat_history: Optional[List[Dict]] = None,
        policy_ids: Optional[List[str]] = None
    ) -> Dict:
        """
        Process a user query through the multi-agent graph.
        
        Args:
            query: User query string
            user_id: Optional user ID
            chat_history: Optional chat history
            policy_ids: Optional list of policy document IDs to search
            
        Returns:
            Dictionary with answer, agent, sources, and metadata
        """
        print(f"\n{'='*60}")
        print(f"[Multi-Agent Graph] Processing query: {query[:50]}...")
        if policy_ids:
            print(f"[Multi-Agent Graph] With {len(policy_ids)} selected policies")
        print(f"{'='*60}\n")
        
        try:
            # Initialize state
            initial_state: AgentState = {
                "query": query,
                "user_id": user_id,
                "chat_history": chat_history or [],
                "policy_ids": policy_ids,
                "agent_type": None,
                "routing_reasoning": None,
                "context_chunks": None,
                "search_results": None,
                "answer": None,
                "sources": None,
                "policy_names": None,  # Initialize policy_names
                "metadata": None,
                "error": None
            }
            
            # Execute the graph
            final_state = self.graph.invoke(initial_state)
            
            # Debug: print what's in final_state
            print(f"[Multi-Agent Graph] Final state keys: {final_state.keys()}")
            print(f"[Multi-Agent Graph] Final state policy_names: {final_state.get('policy_names', 'KEY NOT FOUND')}")
            
            # Extract response
            result = {
                "answer": final_state.get("answer", "I couldn't generate a response."),
                "agent": final_state.get("metadata", {}).get("agent", "unknown"),
                "sources": final_state.get("sources", []),
                "routing_reasoning": final_state.get("routing_reasoning", ""),
                "metadata": final_state.get("metadata", {}),
                "policy_names": final_state.get("policy_names", [])  # Add policy_names!
            }
            
            print(f"[Multi-Agent Graph] Returning policy_names: {result.get('policy_names', [])}")
            print(f"\n[Multi-Agent Graph] Completed - Agent: {result['agent']}\n")
            
            return result
            
        except Exception as e:
            print(f"[Multi-Agent Graph] Error: {e}")
            return {
                "answer": "I apologize, but I encountered an error processing your request. Please try again.",
                "agent": "error",
                "sources": [],
                "routing_reasoning": f"Error: {str(e)}",
                "metadata": {"error": str(e)}
            }


# Global multi-agent graph instance
multi_agent_graph = MultiAgentGraph()