hanzo-mcp 0.7.6__py3-none-any.whl → 0.8.0__py3-none-any.whl

hanzo_mcp/tools/agent/agent.py

@@ -4,42 +4,46 @@ This module provides the AgentTool for delegating tasks to sub-agents,
 supporting both one-off and long-running RPC modes, including A2A communication.
 """
 
-import asyncio
-import json
 import re
+import json
 import time
 import uuid
-from typing import Annotated, TypedDict, Unpack, final, override, Optional, Dict, Any, List
-from collections.abc import Iterable
+import asyncio
 
 # Import litellm with warnings suppressed
 import warnings
+from typing import (
+    Any,
+    Dict,
+    List,
+    Unpack,
+    Optional,
+    Annotated,
+    TypedDict,
+    final,
+    override,
+)
+
 with warnings.catch_warnings():
     warnings.simplefilter("ignore", DeprecationWarning)
-    import litellm
-from mcp.server.fastmcp import Context as MCPContext
-from openai.types.chat import ChatCompletionMessageParam, ChatCompletionToolParam
 from pydantic import Field
+from openai.types.chat import ChatCompletionMessageParam
+from mcp.server.fastmcp import Context as MCPContext
 
+from hanzo_mcp.tools.jupyter import get_read_only_jupyter_tools
+from hanzo_mcp.tools.filesystem import get_read_only_filesystem_tools
+from hanzo_mcp.tools.common.base import BaseTool
 from hanzo_mcp.tools.agent.prompt import (
-    get_allowed_agent_tools,
     get_default_model,
-    get_model_parameters,
     get_system_prompt,
+    get_allowed_agent_tools,
 )
-from hanzo_mcp.tools.agent.tool_adapter import (
-    convert_tools_to_openai_functions,
-)
-from hanzo_mcp.tools.common.base import BaseTool
-from hanzo_mcp.tools.common.batch_tool import BatchTool
 from hanzo_mcp.tools.common.context import (
     ToolContext,
     create_tool_context,
 )
+from hanzo_mcp.tools.common.batch_tool import BatchTool
 from hanzo_mcp.tools.common.permissions import PermissionManager
-from hanzo_mcp.tools.filesystem import get_read_only_filesystem_tools
-from hanzo_mcp.tools.jupyter import get_read_only_jupyter_tools
-
 
 # Parameter types
 Action = Annotated[
@@ -101,6 +105,7 @@ Model = Annotated[
 
 class AgentParams(TypedDict, total=False):
     """Parameters for agent tool."""
+
     action: str
     prompts: Optional[str | List[str]]
     mode: str
@@ -112,8 +117,10 @@ class AgentParams(TypedDict, total=False):
 
 class RPCAgent:
     """Long-running RPC agent."""
-    
-    def __init__(self, agent_id: str, model: str, system_prompt: str, tools: List[BaseTool]):
+
+    def __init__(
+        self, agent_id: str, model: str, system_prompt: str, tools: List[BaseTool]
+    ):
         self.agent_id = agent_id
         self.model = model
         self.system_prompt = system_prompt
@@ -124,12 +131,14 @@ class RPCAgent:
         self.created_at = time.time()
         self.last_used = time.time()
         self.call_count = 0
-        
-    async def call_method(self, method: str, args: Dict[str, Any], tool_ctx: ToolContext) -> str:
+
+    async def call_method(
+        self, method: str, args: Dict[str, Any], tool_ctx: ToolContext
+    ) -> str:
         """Call a method on the RPC agent."""
         self.last_used = time.time()
         self.call_count += 1
-        
+
         # Build prompt based on method
         if method == "search":
             prompt = f"Search for: {args.get('query', 'unknown')}"
@@ -140,22 +149,22 @@ class RPCAgent:
         else:
             # Generic method call
             prompt = f"Method: {method}, Args: {json.dumps(args)}"
-            
+
         # Add to conversation
         self.messages.append({"role": "user", "content": prompt})
-        
+
         # Get response
         # (simplified - would integrate with full agent execution logic)
         response = f"Executed {method} with args {args}"
         self.messages.append({"role": "assistant", "content": response})
-        
+
         return response
 
 
 @final
 class AgentTool(BaseTool):
     """Unified agent tool with one-off and RPC modes."""
-    
+
     def __init__(
         self,
         permission_manager: PermissionManager,
@@ -174,10 +183,10 @@ class AgentTool(BaseTool):
         self.max_tokens_override = max_tokens
         self.max_iterations = max_iterations
         self.max_tool_uses = max_tool_uses
-        
+
         # RPC agent registry
         self._rpc_agents: Dict[str, RPCAgent] = {}
-        
+
         # Available tools
         self.available_tools: list[BaseTool] = []
         self.available_tools.extend(
@@ -201,8 +210,8 @@ class AgentTool(BaseTool):
     def description(self) -> str:
         """Get the tool description."""
         tools = [t.name for t in self.available_tools]
-        
-        return f"""AI agents with tools: {', '.join(tools)}. Actions: run (default), start, call, stop, list.
+
+        return f"""AI agents with tools: {", ".join(tools)}. Actions: run (default), start, call, stop, list.
 
 Usage:
 agent "Search for config files in /project"
@@ -223,10 +232,10 @@ Modes:
         """Execute agent operation."""
         tool_ctx = create_tool_context(ctx)
         await tool_ctx.set_tool_info(self.name)
-        
+
         # Extract action
         action = params.get("action", "run")
-        
+
         # Route to appropriate handler
         if action == "run":
             return await self._handle_run(params, tool_ctx)
@@ -246,30 +255,34 @@ Modes:
         prompts = params.get("prompts")
         if not prompts:
             return "Error: prompts required for run action"
-            
+
         # Convert to list
         if isinstance(prompts, str):
             prompt_list = [prompts]
         else:
             prompt_list = prompts
-            
+
         # Validate prompts
         for prompt in prompt_list:
             if not self._validate_prompt(prompt):
                 return f"Error: Prompt must contain absolute paths starting with /: {prompt[:50]}..."
-        
+
         # Execute agents
         start_time = time.time()
-        
+
         if len(prompt_list) == 1:
             await tool_ctx.info("Launching agent")
-            result = await self._execute_agent(prompt_list[0], params.get("model"), tool_ctx)
+            result = await self._execute_agent(
+                prompt_list[0], params.get("model"), tool_ctx
+            )
         else:
             await tool_ctx.info(f"Launching {len(prompt_list)} agents in parallel")
-            result = await self._execute_multiple_agents(prompt_list, params.get("model"), tool_ctx)
-            
+            result = await self._execute_multiple_agents(
+                prompt_list, params.get("model"), tool_ctx
+            )
+
         execution_time = time.time() - start_time
-        
+
         return f"""Agent execution completed in {execution_time:.2f} seconds.
 
 AGENT RESPONSE:
@@ -280,31 +293,31 @@ AGENT RESPONSE:
         mode = params.get("mode", "oneoff")
         if mode != "rpc":
             return "Error: start action only valid for rpc mode"
-            
+
         # Generate agent ID
         agent_id = str(uuid.uuid4())[:8]
-        
+
         # Get model
         model = params.get("model") or get_default_model(self.model_override)
-        
+
         # Get available tools
         agent_tools = get_allowed_agent_tools(
             self.available_tools,
             self.permission_manager,
         )
-        
+
         # Create system prompt
         system_prompt = get_system_prompt(
             agent_tools,
             self.permission_manager,
         )
-        
+
         # Create RPC agent
         agent = RPCAgent(agent_id, model, system_prompt, agent_tools)
         self._rpc_agents[agent_id] = agent
-        
+
         await tool_ctx.info(f"Started RPC agent {agent_id} with model {model}")
-        
+
         return f"""Started RPC agent:
 - ID: {agent_id}
 - Model: {model}
@@ -317,20 +330,20 @@ Use 'agent --action call --agent-id {agent_id} --method <method> --args <args>'
         agent_id = params.get("agent_id")
         if not agent_id:
             return "Error: agent_id required for call action"
-            
+
         if agent_id not in self._rpc_agents:
             return f"Error: Agent {agent_id} not found. Use 'agent --action list' to see active agents."
-            
+
         method = params.get("method")
         if not method:
             return "Error: method required for call action"
-            
+
         args = params.get("args", {})
-        
+
         # Call agent method
         agent = self._rpc_agents[agent_id]
         await tool_ctx.info(f"Calling {method} on agent {agent_id}")
-        
+
         try:
             result = await agent.call_method(method, args, tool_ctx)
             return f"Agent {agent_id} response:\n{result}"
@@ -343,13 +356,13 @@ Use 'agent --action call --agent-id {agent_id} --method <method> --args <args>'
         agent_id = params.get("agent_id")
         if not agent_id:
             return "Error: agent_id required for stop action"
-            
+
         if agent_id not in self._rpc_agents:
             return f"Error: Agent {agent_id} not found"
-            
+
         agent = self._rpc_agents.pop(agent_id)
         await tool_ctx.info(f"Stopped agent {agent_id}")
-        
+
         return f"""Stopped agent {agent_id}:
 - Runtime: {time.time() - agent.created_at:.2f} seconds
 - Calls: {agent.call_count}"""
@@ -358,7 +371,7 @@ Use 'agent --action call --agent-id {agent_id} --method <method> --args <args>'
         """List active RPC agents."""
         if not self._rpc_agents:
             return "No active RPC agents"
-            
+
         output = ["=== Active RPC Agents ==="]
         for agent_id, agent in self._rpc_agents.items():
             runtime = time.time() - agent.created_at
@@ -368,7 +381,7 @@ Use 'agent --action call --agent-id {agent_id} --method <method> --args <args>'
             output.append(f"  Runtime: {runtime:.2f}s")
             output.append(f"  Idle: {idle:.2f}s")
             output.append(f"  Calls: {agent.call_count}")
-            
+
         return "\n".join(output)
 
     def _validate_prompt(self, prompt: str) -> bool:
@@ -376,30 +389,34 @@ Use 'agent --action call --agent-id {agent_id} --method <method> --args <args>'
         absolute_path_pattern = r"/(?:[^/\s]+/)*[^/\s]+"
         return bool(re.search(absolute_path_pattern, prompt))
 
-    async def _execute_agent(self, prompt: str, model: Optional[str], tool_ctx: ToolContext) -> str:
+    async def _execute_agent(
+        self, prompt: str, model: Optional[str], tool_ctx: ToolContext
+    ) -> str:
         """Execute a single agent (simplified - would use full logic from agent_tool.py)."""
         # This would integrate the full agent execution logic from agent_tool.py
         # For now, return a placeholder
         return f"Executed agent with prompt: {prompt[:100]}..."
 
-    async def _execute_multiple_agents(self, prompts: List[str], model: Optional[str], tool_ctx: ToolContext) -> str:
+    async def _execute_multiple_agents(
+        self, prompts: List[str], model: Optional[str], tool_ctx: ToolContext
+    ) -> str:
         """Execute multiple agents in parallel."""
         tasks = []
         for prompt in prompts:
             task = self._execute_agent(prompt, model, tool_ctx)
             tasks.append(task)
-            
+
         results = await asyncio.gather(*tasks, return_exceptions=True)
-        
+
         formatted_results = []
         for i, result in enumerate(results):
             if isinstance(result, Exception):
-                formatted_results.append(f"Agent {i+1} Error:\n{str(result)}")
+                formatted_results.append(f"Agent {i + 1} Error:\n{str(result)}")
             else:
-                formatted_results.append(f"Agent {i+1} Result:\n{result}")
-                
+                formatted_results.append(f"Agent {i + 1} Result:\n{result}")
+
         return "\n\n---\n\n".join(formatted_results)
 
     def register(self, mcp_server) -> None:
         """Register this tool with the MCP server."""
-        pass
+        pass