meta agent for doc

agent_api.py

@@ -25,7 +25,6 @@ from structured_outputs.api_models import (
 )
 
 from langgraph_agent import CyberLegalAgent
-from utils.conversation_manager import ConversationManager
 from utils.utils import validate_query
 from utils.lightrag_client import LightRAGClient
 from utils import tools
@@ -34,6 +33,7 @@ from subagents.lawyer_messenger import LawyerMessengerAgent
 from prompts.main import SYSTEM_PROMPT_CLIENT, SYSTEM_PROMPT_LAWYER
 from subagents.pdf_analyzer import PDFAnalyzerAgent
 from subagents.doc_editor import DocumentEditorAgent
+from subagents.doc_assistant import DocAssistant
 from langchain_openai import ChatOpenAI
 from mistralai import Mistral
 from langchain_google_genai import ChatGoogleGenerativeAI
@@ -181,7 +181,14 @@ class CyberLegalAPI:
             llm=self.llm_config.openai_llm,
             llm_tool_calling=self.llm_config.gemini_llm
         )
-        self.conversation_manager = ConversationManager()
+        tools.doc_editor_agent = self.doc_editor
+        logger.info("✅ doc_editor_agent initialized globally")
+        # Initialize doc_assistant with Gemini for tool calling
+        self.doc_assistant = DocAssistant(
+            llm=self.llm_config.gemini_llm,
+            tools=tools.tools_for_doc,
+            tools_facade=tools.tools_for_doc_facade
+        )
         logger.info(f"🔧 CyberLegalAPI initialized with {llm_provider.upper()} provider")
     
     def _format_documents_tree(self, node: TreeNode, indent: int = 0) -> str:
@@ -505,16 +512,15 @@ class CyberLegalAPI:
             else:
                 logger.info("ℹ️ No conversation history provided")
             
-            # Call document editor agent
+            # Call doc_assistant (router agent that decides to respond or edit)
             logger.info("=" * 80)
-            logger.info("🤖 CALLING DOCUMENT EDITOR AGENT")
+            logger.info("🤖 CALLING DOC ASSISTANT")
             logger.info("=" * 80)
-            result = await self.doc_editor.edit_document(
+            result = await self.doc_assistant.process_request(
                 doc_text=doc_text,
                 user_instruction=request.instruction,
                 doc_summaries=doc_summaries,
                 conversation_history=conversation_history,
-                max_iterations=10,
                 document_id=request.documentId,
                 user_id=request.userId
             )
@@ -524,18 +530,17 @@ class CyberLegalAPI:
             
             # Log results
             logger.info("=" * 80)
-            logger.info("📊 DOCUMENT EDITING RESULTS")
+            logger.info("📊 DOC ASSISTANT RESULTS")
             logger.info("=" * 80)
             logger.info(f"✅ Success: {result['success']}")
-            logger.info(f"🔄 Iterations: {result['iteration_count']}")
             logger.info(f"⏱️ Processing time: {processing_time:.2f}s")
             logger.info(f"💬 Message: {result['message'][:100]}...")
             
-            # Prepare response - return HTML directly
-            modified_document = result['doc_text'] if result['success'] else None
-            if result['success']:
-                logger.info(f"📏 Modified document size: {len(result['doc_text'])} bytes")
-                logger.info(f"📈 Size change: {len(result['doc_text']) - len(doc_text):+d} bytes")
+            # Prepare response - doc_assistant returns modified_document or None
+            modified_document = result.get('modified_document')
+            if modified_document:
+                logger.info(f"📏 Modified document size: {len(modified_document)} bytes")
+                logger.info(f"📈 Size change: {len(modified_document) - len(doc_text):+d} bytes")
             
             response = DocCreatorResponse(
                 response=result['message'],

agent_states/doc_assistant_state.py

@@ -0,0 +1,37 @@
+#!/usr/bin/env python3
+"""
+State for DocCreatorRouterAgent
+Simple state for the router that decides whether to edit or just respond
+"""
+
+from typing import TypedDict, List, Dict, Any, Optional
+from langchain_core.messages import BaseMessage
+
+
+class DocAssistantState(TypedDict):
+    """
+    State for the document creator router agent.
+    
+    This agent decides whether to:
+    1. Just respond to the user (no document changes needed)
+    2. Call the edit_document tool to modify the document
+    
+    Simple single-pass workflow with optional tool calling.
+    """
+    doc_text: str  # The HTML document content
+    user_instruction: str  # User's instruction or question
+    doc_summaries: Optional[List[str]]  # Optional context from other documents
+    conversation_history: Optional[List[Dict[str, str]]]  # Optional conversation history
+    
+    # Agent response
+    message: Optional[str]  # Agent's response to the user
+    
+    # Modified document (if editing was performed)
+    modified_document: Optional[str]  # The modified document, if edit_document was called
+    
+    # Tool execution results
+    intermediate_steps: List[BaseMessage]  # Conversation history including tool calls
+    
+    # Metadata
+    document_id: Optional[str]  # Optional UUID of the document for live updates
+    user_id: Optional[str]  # Optional user ID for authentication

agent_states/doc_editor_state.py

@@ -19,6 +19,9 @@ class DocEditorState(TypedDict):
     # User instruction
     user_instruction: str
     
+    # Execution plan (optional, provided by DocAssistant)
+    plan: Optional[str]
+    
     # Workflow control
     iteration_count: int
     max_iterations: int

prompts/doc_assistant.py

@@ -0,0 +1,196 @@
+#!/usr/bin/env python3
+"""
+System prompts for the doc creator router agent
+"""
+
+ROUTER_SYSTEM_PROMPT = """You are a Document Router Agent that decides whether to respond to a user's question or modify their HTML document.
+
+## CRITICAL RULES
+
+1. **DECIDE FIRST**: Determine if the user wants to UNDERSTAND the document or MODIFY it.
+
+2. **UNDERSTAND = RESPOND**: If the user asks questions, wants explanations, or requests information → Just respond with your answer. No tool call needed.
+
+3. **MODIFY = USE TOOL**: If the user wants to change, add, remove, or restructure content → Use the edit_document tool.
+
+4. **DOCUMENT IS IN CONTEXT**: The current HTML document is included in this message. You can reference it directly.
+
+5. **ONE DECISION**: Make a clear decision. Either respond or call a tool. Don't hesitate.
+
+6. **BE PRECISE**: If you use edit_document, pass the user's exact instruction.
+
+7. **ACCESS OTHER DOCUMENTS**: You can use retrieve_lawyer_document if you need to reference other documents from the lawyer's database.
+
+## WHEN TO JUST RESPOND
+
+Respond directly without tool calls when the user wants to:
+
+- **Understand content**: "What does this contract say about duration?"
+- **Explain sections**: "Explain Article 3 in simple terms"
+- **Extract information**: "What are the key obligations in this agreement?"
+- **Summarize**: "Give me a summary of this document"
+- **Clarify meaning**: "What does this clause mean?"
+- **Answer questions**: "Is there a termination clause?"
+- **Compare terms**: "What's the difference between Article 1 and Article 2?"
+
+## WHEN TO CALL edit_document
+
+Use the edit_document tool when the user wants to:
+
+- **Change text**: "Change 12 months to 24 months"
+- **Add content**: "Add a GDPR compliance clause"
+- **Remove content**: "Delete the section about confidentiality"
+- **Restructure**: "Move Article 3 after Article 1"
+- **Insert sections**: "Add Article 4 about pricing after Article 3"
+- **Modify clauses**: "Update the liability clause to cap at $10,000"
+- **Update values**: "Change the price to $500"
+- **Rephrase**: "Rewrite this paragraph to be more formal"
+
+## AVAILABLE TOOLS
+
+### edit_document
+Modify the HTML document by providing a detailed execution plan. The instruction will be provided automatically when the tool is triggered.
+
+Parameters:
+- plan: A numbered list of up to 6 steps describing exactly how to execute the modification
+- doc_summaries: Optional context from other documents (for the editor agent)
+- conversation_history: Optional conversation history (for the editor agent)
+
+The plan should be specific and actionable, describing exactly what HTML elements to target and what changes to make.
+
+Example:
+```json
+{
+  "type": "tool_call",
+  "name": "edit_document",
+  "arguments": {
+    "plan": "1. Find the paragraph containing '12 months'\n2. Replace '12 months' with '24 months'\n3. Verify the change is correct"
+  }
+}
+```
+
+Complex example with multiple steps:
+```json
+{
+  "type": "tool_call",
+  "name": "edit_document",
+  "arguments": {
+    "plan": "1. Locate the Article 3 heading\n2. Add a new heading 'Article 4 - GDPR Compliance' after Article 3\n3. Add a GDPR clause paragraph under the new heading\n4. Find the data processing section in the document\n5. Update the data processing paragraph to include GDPR references\n6. Verify all changes maintain valid HTML structure"
+  }
+}
+```
+
+### retrieve_lawyer_document
+Retrieve a specific document from the lawyer's document database for additional context.
+
+Parameters:
+- user_id: The lawyer's user ID
+- file_path: Path to the document (e.g., "Contracts/bail-commercial.pdf")
+
+Example:
+```json
+{
+  "type": "tool_call",
+  "name": "retrieve_lawyer_document",
+  "arguments": {
+    "user_id": "user-uuid",
+    "file_path": "Contracts/contract-template.pdf"
+  }
+}
+```
+
+## DECISION PROCESS
+
+1. **Analyze the user's request**
+2. **Ask yourself**: Does this require CHANGING the document, or just UNDERSTANDING it?
+3. **If understanding → Respond directly**
+4. **If modifying → Call edit_document with the exact instruction**
+5. **If you need other documents → Call retrieve_lawyer_document first, then proceed**
+
+## WORKFLOW EXAMPLES
+
+### Example 1: Understanding (Just Respond)
+User: "What does Article 2 say about termination?"
+→ Analyze: User wants to understand, not modify
+→ Respond: "Article 2 states that either party may terminate with 30 days notice..."
+
+### Example 2: Modification (Use Tool)
+User: "Change 12 months to 24 months"
+→ Analyze: User wants to modify the document
+→ Call edit_document with instruction: "Change 12 months to 24 months"
+
+### Example 3: Complex Modification
+User: "Add a GDPR compliance clause after Article 3"
+→ Analyze: User wants to add content
+→ Call edit_document with instruction: "Add a GDPR compliance clause after Article 3"
+
+### Example 4: Question + Modification
+User: "Is there a penalty clause? If not, add one."
+→ Analyze: Two parts - understand then modify
+→ Respond: "No, there is no penalty clause. I'll add one for you."
+→ Call edit_document with instruction: "Add a penalty clause"
+
+### Example 5: Cross-Document Reference
+User: "Compare this contract's termination clause with the one in Contracts/termination-standard.pdf"
+→ Analyze: Need to reference another document
+→ Call retrieve_lawyer_document for the other document
+→ Compare and respond
+
+## IMPORTANT NOTES
+
+- **Document is provided**: The HTML document is in the message. Read it carefully.
+- **Context available**: Document summaries and conversation history are provided if available.
+- **No iteration needed**: You make one decision. If you call edit_document, it handles the full editing process.
+- **Be helpful**: Whether responding or editing, provide clear, useful information.
+- **Ask if unclear**: If the user's request is ambiguous, ask for clarification.
+
+Remember: Your job is to be efficient and helpful. If no changes are needed, just answer. If changes are needed, use the tool.
+"""
+
+
+def get_router_system_prompt() -> str:
+    """Get the system prompt for the doc creator router agent."""
+    return ROUTER_SYSTEM_PROMPT
+
+
+def get_router_user_prompt(doc_text: str, instruction: str, doc_summaries: list = [], conversation_history: list = []) -> str:
+    """
+    Build the user prompt for the router agent.
+    
+    Args:
+        doc_text: The HTML document content
+        instruction: User's instruction or question
+        doc_summaries: Optional list of document summaries for context
+        conversation_history: Optional conversation history
+    
+    Returns:
+        Formatted user prompt
+    """
+    prompt_parts = []
+    
+    # Add document
+    prompt_parts.append("## Current Document (HTML)")
+    prompt_parts.append(doc_text)
+    prompt_parts.append("")
+    
+    # Add document summaries if provided
+    if doc_summaries:
+        prompt_parts.append("## Related Documents (for context)")
+        for summary in doc_summaries:
+            prompt_parts.append(f"- {summary}")
+        prompt_parts.append("")
+    
+    # Add conversation history if provided
+    if conversation_history:
+        prompt_parts.append("## Previous Conversation")
+        for msg in conversation_history:
+            role = msg.get("role", "").capitalize()
+            content = msg.get("content", "")
+            prompt_parts.append(f"{role}: {content}")
+        prompt_parts.append("")
+    
+    # Add user instruction
+    prompt_parts.append("## User Request")
+    prompt_parts.append(instruction)
+    
+    return "\n".join(prompt_parts)

subagents/doc_assistant.py

@@ -0,0 +1,355 @@
+#!/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 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 utils.tools import _edit_document
+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
+        
+        # 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", [])
+        if not intermediate_steps:
+            return "end"
+        
+        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("💬 Direct response (no tool calls)")
+        
+        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']
+            args = tool_call['args'].copy()
+            
+            # Get the tool function directly from self.tools
+            tool_func = next((t for t in self.tools if t.name == tool_name), None)
+            
+            if tool_func:
+                try:
+                    # Special handling for edit_document - inject all necessary parameters
+                    if tool_name == "edit_document":
+                        logger.info("📝 edit_document tool called - invoking doc_editor_agent")
+                        
+                        # Call the real implementation with all parameters
+                        result = await _edit_document(
+                            doc_text=state["doc_text"],
+                            user_instruction=state["user_instruction"],
+                            plan=args.get("plan", ""),
+                            doc_summaries=state.get("doc_summaries", []),
+                            conversation_history=state.get("conversation_history", []),
+                            max_iterations=10,
+                            document_id=state.get("document_id"),
+                            user_id=state.get("user_id")
+                        )
+                        
+                        # Update state with result
+                        if result.get("ok") and result.get("doc_text"):
+                            state["modified_document"] = result["doc_text"]
+                        
+                        # Add result as tool message
+                        intermediate_steps.append(
+                            ToolMessage(
+                                content=result.get("message", "Document editing completed"),
+                                tool_call_id=tool_call['id'],
+                                name=tool_name
+                            )
+                        )
+                        
+                        logger.info(f"✅ edit_document completed: {result.get('success', False)}")
+                    
+                    # Special handling for retrieve_lawyer_document - inject user_id
+                    elif tool_name == "retrieve_lawyer_document":
+                        from utils.tools import _retrieve_lawyer_document
+                        logger.info(f"📄 retrieve_lawyer_document tool called: {args.get('file_path')}")
+                        
+                        # Inject user_id if available
+                        if "user_id" not in args and state.get("user_id"):
+                            args["user_id"] = state["user_id"]
+                        
+                        result = await _retrieve_lawyer_document(**args)
+                        
+                        intermediate_steps.append(
+                            ToolMessage(
+                                content=result,
+                                tool_call_id=tool_call['id'],
+                                name=tool_name
+                            )
+                        )
+                        
+                        logger.info(f"✅ retrieve_lawyer_document completed")
+                    
+                    # Special handling for query_knowledge_graph
+                    elif tool_name == "query_knowledge_graph":
+                        from utils.tools import _query_knowledge_graph
+                        logger.info(f"🔍 query_knowledge_graph tool called: {args.get('query')}")
+                        
+                        result = await _query_knowledge_graph(**args)
+                        
+                        intermediate_steps.append(
+                            ToolMessage(
+                                content=result,
+                                tool_call_id=tool_call['id'],
+                                name=tool_name
+                            )
+                        )
+                        
+                        logger.info(f"✅ query_knowledge_graph completed")
+                    
+                    # Generic tool invocation (shouldn't happen with current tools)
+                    else:
+                        result = await tool_func.ainvoke(args)
+                        intermediate_steps.append(
+                            ToolMessage(
+                                content=str(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)
+            
+            # Determine result
+            intermediate_steps = final_state.get("intermediate_steps", [])
+            
+            # Check if edit_document was called successfully
+            if final_state.get("modified_document"):
+                logger.info("✅ Document was modified")
+                # Find the tool message with the result
+                for msg in intermediate_steps:
+                    if isinstance(msg, ToolMessage) and msg.name == "edit_document":
+                        return {
+                            "message": msg.content,
+                            "modified_document": final_state["modified_document"],
+                            "success": True
+                        }
+            else:
+                # No modification - just a response
+                logger.info("💬 Router responded without editing")
+                # Find the last AIMessage (agent's response)
+                for msg in reversed(intermediate_steps):
+                    if isinstance(msg, AIMessage) and not (hasattr(msg, 'tool_calls') and msg.tool_calls):
+                        return {
+                            "message": msg.content,
+                            "modified_document": None,
+                            "success": True
+                        }
+            
+            # Fallback
+            return {
+                "message": "Processing completed",
+                "modified_document": None,
+                "success": True
+            }
+        
+        except Exception as e:
+            logger.error("=" * 80)
+            logger.error("❌ DOC ASSISTANT FAILED")
+            logger.error("=" * 80)
+            logger.error(f"❌ Error: {str(e)}")
+            logger.error(f"🔍 Traceback: {traceback.format_exc()}")
+            return {
+                "message": f"Error processing request: {str(e)}",
+                "modified_document": None,
+                "success": False
+            }

subagents/doc_editor.py

@@ -165,12 +165,16 @@ class DocumentEditorAgent:
             
             if doc_summaries:
                 context_msg += "\n\nDocument summaries:\n" + "\n".join(f"- {s}" for s in doc_summaries)
-            
+            plan_formatted=f"""## EXECUTION PLAN
+You have been provided with an execution plan. Follow this plan carefully to complete the task:
+{state['plan']}
+Use this plan as your guide for the editing steps to perform."""
             # Add document and instruction
             full_message = (
                 f"Current document (HTML):\n{state['doc_text']}\n\n"
                 f"{context_msg}\n\n"
                 f"Instruction: {state['user_instruction']}"
+                f"{plan_formatted}"
             )
             intermediate_steps.append(HumanMessage(content=full_message))
             logger.info(f"📚 Context: {len(conversation_history)} hist + {len(doc_summaries)} summaries")
@@ -179,8 +183,6 @@ class DocumentEditorAgent:
         response = await self.llm_with_tools.ainvoke(intermediate_steps)
         intermediate_steps.append(response)
         
-        # Increment iteration count for this LLM call (regardless of tool calls)
-        state["iteration_count"] = iteration_count + 1
         state["intermediate_steps"] = intermediate_steps
         return state
     
@@ -291,6 +293,7 @@ Full conversation history (including all tool calls and results):
         self, 
         doc_text: str, 
         user_instruction: str, 
+        plan: Optional[str] = None,
         doc_summaries: List[str] = [],
         conversation_history: List[Dict[str, str]] = [],
         max_iterations: int = 10,
@@ -303,6 +306,7 @@ Full conversation history (including all tool calls and results):
         Args:
             doc_text: HTML document string
             user_instruction: What changes to make to the document
+            plan: Optional execution plan provided by DocAssistant
             doc_summaries: Optional summaries of the document for context
             conversation_history: Optional previous conversation messages for context
             max_iterations: Maximum number of edit iterations (default: 10)
@@ -344,6 +348,7 @@ Full conversation history (including all tool calls and results):
         initial_state = {
             "doc_text": doc_text,
             "doc_summaries": doc_summaries,
+            "plan": plan,
             "conversation_history": conversation_history,
             "user_instruction": user_instruction,
             "iteration_count": 0,

utils/conversation_manager.py

@@ -1,129 +0,0 @@
-#!/usr/bin/env python3
-"""
-Conversation management for the agent
-"""
-
-from typing import Dict, List, Any
-from datetime import datetime
-
-
-class ConversationManager:
-    """
-    Manages conversation history and context
-    """
-    
-    def __init__(self, max_history: int = 10):
-        self.max_history = max_history
-    
-    def add_exchange(self, history: List[Dict[str, str]], user_query: str, agent_response: str) -> List[Dict[str, str]]:
-        """
-        Add a new user-agent exchange to the conversation history
-        """
-        updated_history = history.copy()
-        
-        # Add user message
-        updated_history.append({
-            "role": "user",
-            "content": user_query,
-            "timestamp": datetime.now().isoformat()
-        })
-        
-        # Add agent response
-        updated_history.append({
-            "role": "assistant", 
-            "content": agent_response,
-            "timestamp": datetime.now().isoformat()
-        })
-        
-        # Keep only the last max_history exchanges (pairs)
-        if len(updated_history) > self.max_history * 2:
-            updated_history = updated_history[-self.max_history * 2:]
-        
-        return updated_history
-    
-    def format_for_lightrag(self, history: List[Dict[str, str]]) -> List[Dict[str, str]]:
-        """
-        Format conversation history for LightRAG API
-        """
-        formatted = []
-        for exchange in history:
-            formatted.append({
-                "role": exchange["role"],
-                "content": exchange["content"]
-            })
-        return formatted
-    
-    def get_context_summary(self, history: List[Dict[str, str]]) -> str:
-        """
-        Generate a summary of recent conversation context
-        """
-        if not history:
-            return "No previous conversation context."
-        
-        recent_exchanges = history[-6:]  # Last 3 exchanges
-        context_parts = []
-        
-        for i, exchange in enumerate(recent_exchanges):
-            role = "User" if exchange["role"] == "user" else "Assistant"
-            context_parts.append(f"{role}: {exchange['content']}")
-        
-        return "\n".join(context_parts)
-
-
-class ConversationFormatter:
-    """
-    Format conversation data for different purposes
-    """
-    
-    @staticmethod
-    def build_conversation_history(history: List[Dict[str, str]], max_turns: int = 10) -> List[Dict[str, str]]:
-        """
-        Build conversation history for LightRAG API
-        """
-        if not history:
-            return []
-        
-        # Take last max_turns pairs (user + assistant)
-        recent_history = history[-max_turns*2:]
-        formatted = []
-        
-        for exchange in recent_history:
-            # Handle both Message objects and dictionary formats
-            if hasattr(exchange, 'role'):
-                role = exchange.role
-                content = exchange.content
-            else:
-                role = exchange["role"]
-                content = exchange["content"]
-            
-            formatted.append({
-                "role": role,
-                "content": content
-            })
-        
-        return formatted
-    
-    @staticmethod
-    def create_context_summary(history: List[Dict[str, str]]) -> str:
-        """
-        Create a summary of conversation context
-        """
-        if not history:
-            return "No previous conversation."
-        
-        recent_exchanges = history[-4:]  # Last 2 exchanges
-        context_parts = []
-        
-        for exchange in recent_exchanges:
-            # Handle both Message objects and dictionary formats
-            if hasattr(exchange, 'role'):
-                role = "User" if exchange.role == "user" else "Assistant"
-                content = exchange.content
-            else:
-                role = "User" if exchange["role"] == "user" else "Assistant"
-                content = exchange["content"]
-            
-            content = content[:100] + "..." if len(content) > 100 else content
-            context_parts.append(f"{role}: {content}")
-        
-        return "\n".join(context_parts)

utils/lightrag_client.py

@@ -165,50 +165,6 @@ class LightRAGClient:
         return ref_list
 
 
-class ResponseProcessor:
-    """
-    Process and enhance LightRAG responses
-    """
-    
-    @staticmethod
-    def extract_main_content(response: Dict[str, Any]) -> str:
-        """
-        Extract the main response content
-        """
-        return response.get("response", "No response available.")
-    
-    @staticmethod
-    def format_references(references: List[str]) -> str:
-        """
-        Format reference list for display
-        """
-        if not references:
-            return ""
-        
-        ref_text = "\n\n**📚 References:**\n"
-        for ref in references:
-            ref_text += f"• {ref}\n"
-        
-        return ref_text
-    
-    @staticmethod
-    def extract_key_entities(response: Dict[str, Any]) -> List[str]:
-        """
-        Extract key entities mentioned in the response
-        """
-        # This could be enhanced if LightRAG provides entity information
-        content = response.get("response", "")
-        
-        # Simple entity extraction based on common legal terms
-        legal_entities = []
-        regulations = ["GDPR", "NIS2", "DORA", "CRA", "eIDAS", "Cyber Resilience Act"]
-        
-        for reg in regulations:
-            if reg.lower() in content.lower():
-                legal_entities.append(reg)
-        
-        return list(set(legal_entities))  # Remove duplicates
-
 
 _lightrag_client_cache: Dict[str, LightRAGClient] = {}
 

utils/tools.py

@@ -17,10 +17,10 @@ import resend
 # Global instances - will be initialized in agent_api.py
 lawyer_selector_agent: Optional[LawyerSelectorAgent] = None
 lawyer_messenger_agent: Optional[LawyerMessengerAgent] = None
-doc_editor_agent: Optional[DocumentEditorAgent] = None
 lightrag_client: Optional[LightRAGClient] = None
-tavily_search = None
+tavily_search: Optional[TavilySearch]  = None
 resend_api_key: Optional[str] = None
+doc_editor_agent: Optional[DocumentEditorAgent] = None
 
 @tool
 async def query_knowledge_graph(
@@ -249,96 +249,6 @@ async def _message_lawyer(conversation_history, user_id) -> str:
     except Exception as e:
         return f"Error sending message to lawyer: {str(e)}"
 
-
-@tool
-async def edit_document() -> str:
-    """
-    Edit a TipTap JSON document according to user instructions.
-    
-    This tool uses an AI agent to intelligently modify documents by making
-    precise edits (replace, add, delete) while maintaining valid JSON structure.
-    The agent works iteratively, validating each edit before proceeding.
-    
-    Use this tool when the user wants to modify a document, such as:
-    - Changing text content (e.g., "Change 12 months to 24 months")
-    - Adding new sections or clauses
-    - Removing unwanted content
-    - Restructuring parts of the document
-    
-    Args:
-        doc_text: The TipTap JSON document as a string (must be canonical format)
-        instruction: What changes to make to the document
-        doc_summaries: Optional list of document summaries for context
-    
-    Returns:
-        A summary of what was changed and the modified document
-    """
-    try:
-        if doc_editor_agent is None:
-            raise ValueError("DocumentEditorAgent not initialized. Please initialize it in agent_api.py")
-        raise ValueError("doc_text and instruction not provided - these should be injected by the agent")
-    except Exception as e:
-        return f"Error editing document: {str(e)}"
-
-
-@tool
-async def _edit_document(doc_text: str, instruction: str, doc_summaries: List[str] = []) -> str:
-    """
-    Edit a TipTap JSON document according to user instructions.
-    
-    This tool uses an AI agent to intelligently modify documents by making
-    precise edits (replace, add, delete) while maintaining valid JSON structure.
-    The agent works iteratively, validating each edit before proceeding.
-    
-    Use this tool when the user wants to modify a document, such as:
-    - Changing text content (e.g., "Change 12 months to 24 months")
-    - Adding new sections or clauses
-    - Removing unwanted content
-    - Restructuring parts of the document
-    
-    Args:
-        doc_text: The TipTap JSON document as a string (must be canonical format with sort_keys=True, indent=2)
-        instruction: What changes to make to the document
-        doc_summaries: Optional list of document summaries for context
-    
-    Returns:
-        A summary of what was changed and the modified document
-    """
-    try:
-        if doc_editor_agent is None:
-            raise ValueError("DocumentEditorAgent not initialized. Please initialize it in agent_api.py")
-        
-        # Ensure doc_text is canonical format
-        try:
-            doc_obj = json.loads(doc_text)
-            canonical_doc_text = json.dumps(doc_obj, ensure_ascii=False, sort_keys=True, indent=2)
-        except json.JSONDecodeError:
-            return f"Error: Invalid JSON provided in doc_text"
-        
-        # Run the document editor agent
-        result = await doc_editor_agent.edit_document(
-            doc_text=canonical_doc_text,
-            user_instruction=instruction,
-            doc_summaries=doc_summaries
-        )
-        
-        if result["success"]:
-            output = [
-                "✅ Document editing completed successfully",
-                "=" * 80,
-                f"\nSummary: {result['message']}",
-                f"\nIterations: {result['iteration_count']}",
-                "\nModified document (TipTap JSON):",
-                result["doc_text"]
-            ]
-            return "\n".join(output)
-        else:
-            return f"❌ Failed to edit document: {result['message']}"
-    
-    except Exception as e:
-        return f"Error editing document: {str(e)}"
-
-
 @tool
 async def retrieve_lawyer_document(file_path: str) -> str:
     """
@@ -447,9 +357,93 @@ async def _retrieve_lawyer_document(
         return f"Error retrieving document: {str(e)}"
 
 
+# ============ DOC ASSISTANT TOOLS ============
+
+@tool
+async def edit_document(plan: str) -> str:
+    """
+    Edit the HTML document by providing a detailed execution plan.
+    
+    This tool is a facade - the router agent calls this when it decides that
+    document editing is needed. The actual editing will be performed by the
+    DocumentEditorAgent.
+    
+    Args:
+        plan: A numbered list of up to 6 steps describing exactly how to execute the modification
+    
+    Returns:
+        Placeholder response - actual implementation is in the router
+    
+    Example:
+        plan = "1. Find the paragraph containing '12 months'\n2. Replace '12 months' with '24 months'\n3. Verify the change is correct"
+    """
+    return
+
+
+@tool
+async def _edit_document(
+    doc_text: str,
+    user_instruction: str,
+    plan: str,
+    doc_summaries: list = [],
+    conversation_history: list = [],
+    max_iterations: int = 10,
+    document_id: str = None,
+    user_id: str = None
+) -> str:
+    """
+    Edit the HTML document by launching the DocumentEditorAgent workflow.
+    
+    This is the real implementation that actually executes the document editor
+    graph with all necessary parameters.
+    
+    Args:
+        doc_text: The HTML document content
+        user_instruction: User's instruction for editing
+        plan: A numbered list of up to 6 steps describing exactly how to execute the modification
+        doc_summaries: Optional list of document summaries for context
+        conversation_history: Optional conversation history
+        max_iterations: Maximum number of edit iterations (default: 10)
+        document_id: Optional UUID of the document for live updates
+        user_id: Optional user ID for authentication
+    
+    Returns:
+        Result message from the editor agent
+    """
+    try:
+        if doc_editor_agent is None:
+            raise ValueError("DocumentEditorAgent not initialized. Please initialize it in agent_api.py")
+        
+        # Invoke the document editor agent
+        result = await doc_editor_agent.edit_document(
+            doc_text=doc_text,
+            user_instruction=user_instruction,
+            plan=plan,
+            doc_summaries=doc_summaries,
+            conversation_history=conversation_history,
+            max_iterations=max_iterations,
+            document_id=document_id,
+            user_id=user_id
+        )
+        
+        # Return the result as a string
+        if result.get("success"):
+            return f"Document edited successfully: {result.get('message', 'Completed')}"
+        else:
+            return f"Document editing failed: {result.get('message', 'Unknown error')}"
+    
+    except Exception as e:
+        return f"Error editing document: {str(e)}"
+
+
 # Export tool sets for different user types
-tools_for_client_facade = [query_knowledge_graph, find_lawyers, message_lawyer, search_web, edit_document]
-tools_for_client = [_query_knowledge_graph, _find_lawyers, _message_lawyer, search_web, _edit_document]
-tools_for_lawyer_facade = [query_knowledge_graph, search_web, edit_document, retrieve_lawyer_document]
-tools_for_lawyer = [_query_knowledge_graph, search_web, _edit_document, _retrieve_lawyer_document]
+tools_for_client_facade = [query_knowledge_graph, find_lawyers, message_lawyer, search_web]
+tools_for_client = [_query_knowledge_graph, _find_lawyers, _message_lawyer, search_web ]
+tools_for_lawyer_facade = [query_knowledge_graph, search_web, retrieve_lawyer_document]
+tools_for_lawyer = [_query_knowledge_graph, search_web, _retrieve_lawyer_document]
+
+# Tools for DocAssistant (document router)
+tools_for_doc_facade = [query_knowledge_graph, retrieve_lawyer_document, edit_document]
+tools_for_doc = [_query_knowledge_graph, _retrieve_lawyer_document, _edit_document]
+
 tools = tools_for_client