fix tool calling chat_agent et doc_assistant

agents/chat_agent.py

@@ -105,27 +105,28 @@ class CyberLegalAgent:
                     args["conversation_history"] = state.get("conversation_history", [])
                     logger.info(f"📝 Injecting conversation_history to {tool_call['name']}: {len(args['conversation_history'])} messages")
                 
-                    # Inject jurisdiction for query_knowledge_graph tool
-                    if tool_call['name'] == "query_knowledge_graph":
-                        args["jurisdiction"] = state.get("jurisdiction")
-                        logger.info(f"🌍 Injecting jurisdiction: {args['jurisdiction']}")
-                    
-                    # Inject user_id for message_lawyer tool
-                    if tool_call['name'] == "message_lawyer":
-                        args["user_id"] = state.get("user_id")
-                        logger.info(f"👤 Injecting user_id: {args['user_id']}")
-                    
-                    # Inject user_id for retrieve_lawyer_document tool
-                    if tool_call['name'] == "retrieve_lawyer_document":
-                        args["user_id"] = state.get("user_id")
-                        logger.info(f"📄 Injecting user_id for retrieve_lawyer_document: {args['user_id']}")
-                    
-                    # Inject user_id for create_draft_document tool
-                    if tool_call['name'] == "create_draft_document":
-                        args["user_id"] = state.get("user_id")
-                        logger.info(f"📝 Injecting user_id for create_draft_document: {args['user_id']}")
-                    
-                    tool_call['name']="_" + tool_call['name']
+                # Inject jurisdiction for query_knowledge_graph tool
+                if tool_call['name'] == "query_knowledge_graph":
+                    args["jurisdiction"] = state.get("jurisdiction")
+                    logger.info(f"🌍 Injecting jurisdiction: {args['jurisdiction']}")
+                
+                # Inject user_id for message_lawyer tool
+                if tool_call['name'] == "message_lawyer":
+                    args["user_id"] = state.get("user_id")
+                    logger.info(f"👤 Injecting user_id: {args['user_id']}")
+                
+                # Inject user_id for retrieve_lawyer_document tool
+                if tool_call['name'] == "retrieve_lawyer_document":
+                    args["user_id"] = state.get("user_id")
+                    logger.info(f"📄 Injecting user_id for retrieve_lawyer_document: {args['user_id']}")
+                
+                # Inject user_id for create_draft_document tool
+                if tool_call['name'] == "create_draft_document":
+                    args["user_id"] = state.get("user_id")
+                    logger.info(f"📝 Injecting user_id for create_draft_document: {args['user_id']}")
+                
+                # Convert to implementation name
+                tool_call['name'] = "_" + tool_call['name']
                 
                 result = await tool_func.ainvoke(args)
                 logger.info(f"🔧 Tool {tool_call} returned: {result}")

agents/doc_assistant.py

@@ -132,7 +132,6 @@ class DocAssistant:
         
         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 (add underscore prefix)
             tool_func = next((t for t in self.tools if t.name == "_" + tool_name), None)
@@ -140,6 +139,8 @@ class DocAssistant:
             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")
                         
@@ -150,26 +151,27 @@ class DocAssistant:
                         args["max_iterations"] = 10
                         args["document_id"] = state.get("document_id")
                         args["user_id"] = state.get("user_id")
-                        result = await tool_func.ainvoke(args)
-                        doc_text=result['doc_text']
-                        message=f"Document was edited with this summary :{result['final_summary']}"
-                        state['modified_document']=doc_text
-                        state['message']=message
-                        
-                        logger.info(f"✅ edit_document called - ending router workflow")
                     
                     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"]
-                        
-                        
-                    logger.info(f"Lauching tool:{tool_name} with args {json.dumps(args)}")
+                            
                     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=message,
+                            content=tool_result,
                             tool_call_id=tool_call['id'],
                             name=tool_name
                         )