agents/doc_assistant.py

#!/usr/bin/env python3
"""
DocAssistant - LangGraph agent that routes between responding and editing
Single node agent that decides whether to respond to user questions or edit the document
"""
import json
import logging
import traceback
from typing import Dict, Any, List, Optional
from langgraph.graph import StateGraph, END
from langchain_core.messages import SystemMessage, HumanMessage, AIMessage, ToolMessage

from agent_states.doc_assistant_state import DocAssistantState
from prompts.doc_assistant import get_router_system_prompt, get_router_user_prompt

logger = logging.getLogger(__name__)


class DocAssistant:
    """
    Router agent that decides whether to respond to user questions or edit the document.
    
    Workflow:
    - Single node with optional tool calling
    - Document is in user prompt (no inspect needed)
    - Agent decides: respond directly OR call edit_document tool
    - No looping at router level (single pass)
    """
    
    def __init__(self, llm, tools: List[Any] = None, tools_facade: List[Any] = None):
        """
        Initialize the router agent.
        
        Args:
            llm: LLM for routing decisions (with tool calling capability)
            tools: Real tool implementations for lookup
            tools_facade: Facade tools for LLM (minimal parameters)
        """
        self.tools = tools
        self.tools_facade = tools_facade if tools_facade else tools
        self.llm= llm
        # Bind facade tools to LLM - optional tool calling
        self.llm_with_tools = self.llm.bind_tools(self.tools_facade)
        
        logger.info("🔧 DocAssistant initialized")
        logger.info(f"🤖 Using {type(llm).__name__} for routing decisions")
        logger.info(f"🛠️ Tools available: {[t.name for t in self.tools_facade]}")
        
        self.workflow = self._build_workflow()
    
    def _build_workflow(self) -> StateGraph:
        """Build the LangGraph workflow for the router agent."""
        workflow = StateGraph(DocAssistantState)
        workflow.add_node("agent", self._agent_node)
        workflow.add_node("tools", self._tools_node)
        workflow.set_entry_point("agent")
        
        # Conditional edge after agent: go to tools if tool calls, else END
        workflow.add_conditional_edges(
            "agent",
            self._should_call_tools,
            {"tools": "tools", "end": END}
        )
        
        # Conditional edge after tools: END if edit_document was called, else continue
        workflow.add_conditional_edges(
            "tools",
            self._after_tools,
            {"end": END, "continue": "agent"}
        )
        
        return workflow.compile()
    
    def _should_call_tools(self, state: DocAssistantState) -> str:
        """
        Decide whether to call tools after agent node.
        
        Returns:
            "tools" if agent made tool calls, "end" otherwise
        """
        intermediate_steps = state.get("intermediate_steps", [])
        last_message = intermediate_steps[-1]
        
        # Check if last message has tool calls
        if hasattr(last_message, 'tool_calls') and last_message.tool_calls:
            tool_names = [tc['name'] for tc in last_message.tool_calls]
            logger.info(f"🔧 Agent calling tools: {tool_names}")
            return "tools"
        
        return "end"
    
    async def _agent_node(self, state: DocAssistantState) -> DocAssistantState:
        """Agent node: Generate response or tool call based on user request."""
        intermediate_steps = state.get("intermediate_steps", [])
        
        logger.info("🎯 Router agent processing request...")
        
        # Build messages
        messages = [
            SystemMessage(content=get_router_system_prompt()),
            HumanMessage(content=get_router_user_prompt(
                doc_text=state["doc_text"],
                instruction=state["user_instruction"],
                doc_summaries=state.get("doc_summaries", []),
                conversation_history=state.get("conversation_history", [])
            ))
        ]
        
        logger.info(f"📏 Document size: {len(state['doc_text'])} bytes")
        logger.info(f"📋 Instruction: {state['user_instruction'][:100]}{'...' if len(state['user_instruction']) > 100 else ''}")
        
        # Call LLM
        response = await self.llm_with_tools.ainvoke(messages)
        intermediate_steps.append(response)
        
        logger.info(f"📊 Response type: {type(response).__name__}")
        if hasattr(response, 'tool_calls') and response.tool_calls:
            logger.info(f"🔧 Tool calls: {[tc['name'] for tc in response.tool_calls]}")
        else:
            logger.info(f"💬 Direct response (no tool calls):{response.content}")
            state['message']= response.content
        state["intermediate_steps"] = intermediate_steps
        return state
    
    async def _tools_node(self, state: DocAssistantState) -> DocAssistantState:
        """Tools node: Execute tool calls (edit_document or retrieve_lawyer_document)."""
        intermediate_steps = state.get("intermediate_steps", [])
        last_message = intermediate_steps[-1]
        
        if not (hasattr(last_message, 'tool_calls') and last_message.tool_calls):
            return state
        
        for tool_call in last_message.tool_calls:
            tool_name = tool_call['name']
            
            # Get the tool function directly from self.tools (add underscore prefix)
            tool_func = next((t for t in self.tools if t.name == "_" + tool_name), None)
            
            if tool_func:
                try:
                    args = tool_call['args'].copy()
                    logger.info(f"Launching tool: {tool_name} with args {json.dumps(args, default=str)}")
                    
                    if tool_name == "edit_document":
                        logger.info("📝 edit_document tool called - invoking doc_editor_agent")
                        
                        args["doc_text"] = state["doc_text"]
                        args["user_instruction"] = state["user_instruction"]
                        args["doc_summaries"] = state.get("doc_summaries", [])
                        args["conversation_history"] = state.get("conversation_history", [])
                        args["max_iterations"] = 10
                        args["document_id"] = state.get("document_id")
                        args["user_id"] = state.get("user_id")
                    
                    elif tool_name == "retrieve_lawyer_document":
                        logger.info(f"📄 retrieve_lawyer_document tool called: {args.get('file_path')}")
                        
                        if "user_id" not in args and state.get("user_id"):
                            args["user_id"] = state["user_id"]
                            
                    result = await tool_func.ainvoke(args)
                    
                    if tool_name == "edit_document":
                        doc_text = result['doc_text']
                        tool_result = f"Document was edited with this summary: {result['final_summary']}"
                        state['modified_document'] = doc_text
                        state['message'] = tool_result
                        logger.info(f"✅ edit_document completed - ending router workflow")
                    else:
                        tool_result = result

                    intermediate_steps.append(
                        ToolMessage(
                            content=tool_result,
                            tool_call_id=tool_call['id'],
                            name=tool_name
                        )
                    )
                except Exception as e:
                    logger.error(f"❌ Error executing {tool_name}: {str(e)}")
                    intermediate_steps.append(
                        ToolMessage(
                            content=f"Error: {str(e)}",
                            tool_call_id=tool_call['id'],
                            name=tool_name
                        )
                    )
            else:
                logger.warning(f"⚠️ Tool function not found for {tool_name}")
        
        state["intermediate_steps"] = intermediate_steps
        return state
    
    def _after_tools(self, state: DocAssistantState) -> str:
        """
        Decide whether to continue after tools node.
        
        Returns:
            "end" if edit_document was called (stops workflow),
            "continue" if retrieve_lawyer_document or query_knowledge_graph were called (allows more tool calls or response)
        """
        intermediate_steps = state.get("intermediate_steps", [])
        
        # Check if edit_document was called
        for msg in reversed(intermediate_steps):
            if isinstance(msg, ToolMessage) and msg.name == "edit_document":
                logger.info("✅ edit_document called - ending router workflow")
                return "end"
        
        # If edit_document wasn't called, continue (allows agent to make more tool calls or respond)
        logger.info("🔄 Continuing router workflow (edit_document not yet called)")
        return "continue"
    
    async def process_request(
        self,
        doc_text: str,
        user_instruction: str,
        doc_summaries: List[str] = [],
        conversation_history: List[Dict[str, str]] = [],
        document_id: Optional[str] = None,
        user_id: Optional[str] = None
    ) -> Dict[str, Any]:
        """
        Process the user's request and decide whether to respond or edit.
        
        Args:
            doc_text: The HTML document content
            user_instruction: User's instruction or question
            doc_summaries: Optional list of document summaries for context
            conversation_history: Optional conversation history
            document_id: Optional UUID of the document for live updates
            user_id: Optional user ID for authentication
        
        Returns:
            Dict with:
                - message: Response message to the user
                - modified_document: Modified document (if editing was done) or None
                - success: Boolean indicating success
        """
        logger.info("=" * 80)
        logger.info("🎯 DOC ASSISTANT STARTING")
        logger.info("=" * 80)
        logger.info(f"📏 Document size: {len(doc_text)} bytes")
        logger.info(f"📋 Instruction: {user_instruction[:100]}{'...' if len(user_instruction) > 100 else ''}")
        logger.info(f"📚 Document summaries: {len(doc_summaries)}")
        logger.info(f"💬 Conversation history: {len(conversation_history)} messages")
        
        try:
            # Initialize state
            initial_state = {
                "doc_text": doc_text,
                "doc_summaries": doc_summaries,
                "conversation_history": conversation_history,
                "user_instruction": user_instruction,
                "intermediate_steps": [],
                "document_id": document_id,
                "user_id": user_id,
                "modified_document": None
            }
            
            # Run workflow
            logger.info("🔄 Invoking router workflow...")
            final_state = await self.workflow.ainvoke(initial_state)
            
            modified_doc=final_state.get("modified_document")
            message=final_state.get("message")
            return {
                "message": message,
                "modified_document": modified_doc,
                "success": True
            }
        
        except Exception as e:
            logger.error("=" * 80)
            logger.error("❌ DOC ASSISTANT FAILED")
            logger.error("=" * 80)
            logger.error(f"📍 Location: subagents/doc_assistant.py:{traceback.extract_tb(e.__traceback__)[-1].lineno} (process_request)")
            logger.error("")
            logger.error("📋 Input Parameters:")
            logger.error(f"   - User Instruction: {user_instruction[:100] if len(user_instruction) > 100 else user_instruction}")
            logger.error(f"   - Document Size: {len(doc_text):,} bytes")
            if document_id:
                logger.error(f"   - Document ID: {document_id}")
            if user_id:
                logger.error(f"   - User ID: {user_id}")
            logger.error(f"   - Document Summaries: {len(doc_summaries)}")
            logger.error(f"   - Conversation History: {len(conversation_history)} messages")
            logger.error("")
            logger.error("🤖 LLM Configuration:")
            logger.error(f"   - LLM: {type(self.llm).__name__}")
            logger.error(f"   - Tools Available: {len(self.tools)}")
            logger.error(f"   - Tool Names: {', '.join([t.name for t in self.tools])}")
            logger.error(f"   - Facade Tools: {len(self.tools_facade)}")
            logger.error("")
            logger.error("🔄 Workflow State:")
            if final_state and "intermediate_steps" in final_state:
                intermediate_steps = final_state["intermediate_steps"]
                logger.error(f"   - Intermediate Steps: {len(intermediate_steps)} messages")
                if intermediate_steps:
                    last_msg = intermediate_steps[-1]
                    logger.error(f"   - Last Message Type: {type(last_msg).__name__}")
                    if hasattr(last_msg, 'tool_calls') and last_msg.tool_calls:
                        logger.error(f"   - Last Tool Calls: {', '.join([tc['name'] for tc in last_msg.tool_calls])}")
            logger.error("")
            logger.error("🔍 Exception Details:")
            logger.error(f"   Type: {type(e).__name__}")
            logger.error(f"   Message: {str(e)}")
            logger.error("")
            logger.error("📝 Full Traceback:")
            logger.error(traceback.format_exc())
            logger.error("")
            logger.error("💾 Document Preview (first 200 chars):")
            logger.error(f"   {doc_text[:200]}")
            logger.error("=" * 80)
            return {
                "message": f"Error processing request: {str(e)}",
                "modified_document": None,
                "success": False
            }