hanzo-mcp 0.7.7__py3-none-any.whl → 0.8.1__py3-none-any.whl

hanzo_mcp/tools/agent/swarm_tool.py

@@ -4,49 +4,90 @@ This module implements the SwarmTool that leverages the hanzo-agents SDK
 for sophisticated multi-agent orchestration with flexible network topologies.
 """
 
-import asyncio
-import json
 import os
-from typing import Annotated, Any, TypedDict, Unpack, final, override, Optional, Dict, List
-from pathlib import Path
+from typing import (
+    Any,
+    Dict,
+    List,
+    Unpack,
+    Optional,
+    TypedDict,
+    final,
+    override,
+)
 
 from mcp.server import FastMCP
 from mcp.server.fastmcp import Context as MCPContext
-from pydantic import Field
 
 # Import hanzo-agents SDK with fallback
 try:
     from hanzo_agents import (
-        Agent, State, Network, Tool, History,
-        ModelRegistry, InferenceResult, ToolCall,
+        Tool,
+        Agent,
+        State,
         Router,
+        History,
+        Network,
+        ToolCall,
+        ModelRegistry,
+        InferenceResult,
     )
+
     HANZO_AGENTS_AVAILABLE = True
 except ImportError:
     # Define minimal stubs if hanzo-agents is not available
     HANZO_AGENTS_AVAILABLE = False
-    class Agent: pass
-    class State: pass
-    class Network: pass
-    class Tool: pass
-    class History: pass
-    class ModelRegistry: pass
-    class InferenceResult: pass
-    class ToolCall: pass
-    class Router: pass
+
+    class Agent:
+        pass
+
+    class State:
+        pass
+
+    class Network:
+        pass
+
+    class Tool:
+        pass
+
+    class History:
+        pass
+
+    class ModelRegistry:
+        pass
+
+    class InferenceResult:
+        pass
+
+    class ToolCall:
+        pass
+
+    class Router:
+        pass
+
 
 # Import optional components with fallbacks
 try:
-    from hanzo_agents import DeterministicRouter, LLMRouter, HybridRouter
+    from hanzo_agents import LLMRouter, HybridRouter, DeterministicRouter
 except ImportError:
     try:
         # Try core module import
-        from hanzo_agents.core.router import DeterministicRouter, LLMRouter, HybridRouter
+        from hanzo_agents.core.router import (
+            LLMRouter,
+            HybridRouter,
+            DeterministicRouter,
+        )
     except ImportError:
         # Define stubs if not available
-        class DeterministicRouter: pass
-        class LLMRouter: pass
-        class HybridRouter: pass
+        class DeterministicRouter:
+            pass
+
+        class LLMRouter:
+            pass
+
+        class HybridRouter:
+            pass
+
 
 try:
     from hanzo_agents import create_memory_kv, create_memory_vector
@@ -56,48 +97,69 @@ except ImportError:
         from hanzo_agents.core.memory import create_memory_kv, create_memory_vector
     except ImportError:
         # Define stubs if not available
-        def create_memory_kv(*args, **kwargs): pass
-        def create_memory_vector(*args, **kwargs): pass
+        def create_memory_kv(*args, **kwargs):
+            pass
+
+        def create_memory_vector(*args, **kwargs):
+            pass
+
 
 try:
     from hanzo_agents import sequential_router, conditional_router, state_based_router
 except ImportError:
     try:
         # Try core module import
-        from hanzo_agents.core.router import sequential_router, conditional_router, state_based_router
+        from hanzo_agents.core.router import (
+            sequential_router,
+            conditional_router,
+            state_based_router,
+        )
     except ImportError:
         # Define stubs if not available
-        def sequential_router(*args, **kwargs): pass
-        def conditional_router(*args, **kwargs): pass
-        def state_based_router(*args, **kwargs): pass
+        def sequential_router(*args, **kwargs):
+            pass
+
+        def conditional_router(*args, **kwargs):
+            pass
+
+        def state_based_router(*args, **kwargs):
+            pass
+
 
 try:
     from hanzo_agents.core.cli_agent import (
-        ClaudeCodeAgent, OpenAICodexAgent,
-        GeminiAgent, GrokAgent
+        GrokAgent,
+        GeminiAgent,
+        ClaudeCodeAgent,
+        OpenAICodexAgent,
     )
 except ImportError:
     # Define stub classes if not available
     class ClaudeCodeAgent(Agent):
         pass
+
     class OpenAICodexAgent(Agent):
         pass
+
     class GeminiAgent(Agent):
         pass
+
     class GrokAgent(Agent):
         pass
 
-from hanzo_mcp.tools.agent.agent_tool import MCPAgent, MCPToolAdapter
+
+from hanzo_mcp.tools.jupyter import get_read_only_jupyter_tools
+from hanzo_mcp.tools.filesystem import Edit, MultiEdit, get_read_only_filesystem_tools
 from hanzo_mcp.tools.common.base import BaseTool
 from hanzo_mcp.tools.common.context import create_tool_context
-from hanzo_mcp.tools.common.permissions import PermissionManager
-from hanzo_mcp.tools.filesystem import get_read_only_filesystem_tools, Edit, MultiEdit
-from hanzo_mcp.tools.jupyter import get_read_only_jupyter_tools
+from hanzo_mcp.tools.agent.agent_tool import MCPAgent
 from hanzo_mcp.tools.common.batch_tool import BatchTool
+from hanzo_mcp.tools.common.permissions import PermissionManager
 
 
 class AgentNode(TypedDict):
     """Node in the agent network."""
+
     id: str
     query: str
     model: Optional[str]
@@ -109,6 +171,7 @@ class AgentNode(TypedDict):
 
 class SwarmConfig(TypedDict):
     """Configuration for an agent network."""
+
     agents: Dict[str, AgentNode]
     entry_point: Optional[str]
     topology: Optional[str]
@@ -116,6 +179,7 @@ class SwarmConfig(TypedDict):
 
 class SwarmToolParams(TypedDict):
     """Parameters for the SwarmTool."""
+
     config: SwarmConfig
     query: str
     context: Optional[str]
@@ -126,11 +190,10 @@ class SwarmToolParams(TypedDict):
 
 class SwarmState(State):
     """State for swarm execution."""
-    
-    def __init__(self, 
-                 config: SwarmConfig,
-                 initial_query: str,
-                 context: Optional[str] = None):
+
+    def __init__(
+        self, config: SwarmConfig, initial_query: str, context: Optional[str] = None
+    ):
         """Initialize swarm state."""
         super().__init__()
         self.config = config
@@ -140,28 +203,30 @@ class SwarmState(State):
         self.completed_agents = set()
         self.current_agent = None
         self.execution_order = []
-    
+
     def to_dict(self) -> Dict[str, Any]:
         """Convert to dictionary."""
         base_dict = super().to_dict()
-        base_dict.update({
-            "config": self.config,
-            "initial_query": self.initial_query,
-            "context": self.context,
-            "agent_results": self.agent_results,
-            "completed_agents": list(self.completed_agents),
-            "current_agent": self.current_agent,
-            "execution_order": self.execution_order
-        })
+        base_dict.update(
+            {
+                "config": self.config,
+                "initial_query": self.initial_query,
+                "context": self.context,
+                "agent_results": self.agent_results,
+                "completed_agents": list(self.completed_agents),
+                "current_agent": self.current_agent,
+                "execution_order": self.execution_order,
+            }
+        )
         return base_dict
-    
+
     @classmethod
     def from_dict(cls, data: Dict[str, Any]) -> "SwarmState":
         """Create from dictionary."""
         state = cls(
             config=data.get("config", {}),
             initial_query=data.get("initial_query", ""),
-            context=data.get("context")
+            context=data.get("context"),
         )
         state.agent_results = data.get("agent_results", {})
         state.completed_agents = set(data.get("completed_agents", []))
@@ -172,35 +237,37 @@ class SwarmState(State):
 
 class SwarmAgent(MCPAgent):
     """Agent that executes within a swarm network."""
-    
-    def __init__(self,
-                 agent_id: str,
-                 agent_config: AgentNode,
-                 available_tools: List[BaseTool],
-                 permission_manager: PermissionManager,
-                 ctx: MCPContext,
-                 **kwargs):
+
+    def __init__(
+        self,
+        agent_id: str,
+        agent_config: AgentNode,
+        available_tools: List[BaseTool],
+        permission_manager: PermissionManager,
+        ctx: MCPContext,
+        **kwargs,
+    ):
         """Initialize swarm agent."""
         # Set name and description from config
         self.name = agent_id
         self.description = agent_config.get("role", f"Agent {agent_id}")
         self.agent_config = agent_config
-        
+
         # Initialize with specified model
         model = agent_config.get("model")
         if model:
             model = self._normalize_model(model)
         else:
             model = "model://anthropic/claude-3-5-sonnet-20241022"
-        
+
         super().__init__(
             available_tools=available_tools,
             permission_manager=permission_manager,
             ctx=ctx,
             model=model,
-            **kwargs
+            **kwargs,
         )
-    
+
     def _normalize_model(self, model: str) -> str:
         """Normalize model names to full format."""
         model_map = {
@@ -211,35 +278,37 @@ class SwarmAgent(MCPAgent):
             "gemini-1.5-pro": "model://google/gemini-1.5-pro",
             "gemini-1.5-flash": "model://google/gemini-1.5-flash",
         }
-        
+
         # Check if it's already a model:// URI
         if model.startswith("model://"):
             return model
-        
+
         # Check mapping
         if model in model_map:
             return model_map[model]
-        
+
         # Assume it's a provider/model format
         if "/" in model:
             return f"model://{model}"
-        
+
         # Default to anthropic
         return f"model://anthropic/{model}"
-    
-    async def run(self, state: SwarmState, history: History, network: Network) -> InferenceResult:
+
+    async def run(
+        self, state: SwarmState, history: History, network: Network
+    ) -> InferenceResult:
         """Execute the swarm agent."""
         # Build prompt with context
         prompt_parts = []
-        
+
         # Add role context
         if self.agent_config.get("role"):
             prompt_parts.append(f"Your role: {self.agent_config['role']}")
-        
+
         # Add shared context
         if state.context:
             prompt_parts.append(f"Context:\n{state.context}")
-        
+
         # Add inputs from connected agents
         receives_from = self.agent_config.get("receives_from", [])
         if receives_from:
@@ -247,41 +316,42 @@ class SwarmAgent(MCPAgent):
             for agent_id in receives_from:
                 if agent_id in state.agent_results:
                     inputs[agent_id] = state.agent_results[agent_id]
-            
+
             if inputs:
                 prompt_parts.append("Input from previous agents:")
                 for input_agent, input_result in inputs.items():
                     prompt_parts.append(f"\n--- From {input_agent} ---\n{input_result}")
-        
+
         # Add file context if specified
         if self.agent_config.get("file_path"):
             prompt_parts.append(f"\nFile to work on: {self.agent_config['file_path']}")
-        
+
         # Add the main query
         prompt_parts.append(f"\nTask: {self.agent_config['query']}")
-        
+
         # Add initial query if this is entry point
         if state.current_agent == state.config.get("entry_point"):
             prompt_parts.append(f"\nMain objective: {state.initial_query}")
-        
+
         full_prompt = "\n\n".join(prompt_parts)
-        
+
         # Execute using base class
         messages = [
             {"role": "system", "content": self._get_system_prompt()},
-            {"role": "user", "content": full_prompt}
+            {"role": "user", "content": full_prompt},
         ]
-        
+
         # Call model
         from hanzo_agents import ModelRegistry
+
         adapter = ModelRegistry.get_adapter(self.model)
         response = await adapter.chat(messages)
-        
+
         # Store result in state
         state.agent_results[self.name] = response
         state.completed_agents.add(self.name)
         state.execution_order.append(self.name)
-        
+
         # Return result
         return InferenceResult(
             agent=self.name,
@@ -289,39 +359,39 @@ class SwarmAgent(MCPAgent):
             metadata={
                 "agent_id": self.name,
                 "role": self.agent_config.get("role"),
-                "connections": self.agent_config.get("connections", [])
-            }
+                "connections": self.agent_config.get("connections", []),
+            },
         )
 
 
 class SwarmRouter(DeterministicRouter):
     """Router for swarm agent orchestration."""
-    
+
     def __init__(self, swarm_config: SwarmConfig):
         """Initialize swarm router."""
         self.swarm_config = swarm_config
         self.agents_config = swarm_config["agents"]
         self.entry_point = swarm_config.get("entry_point")
-        
+
         # Build dependency graph
         self.dependencies = {}
         self.dependents = {}
-        
+
         for agent_id, config in self.agents_config.items():
             # Dependencies (agents this one waits for)
             self.dependencies[agent_id] = set(config.get("receives_from", []))
-            
+
             # Dependents (agents that wait for this one)
             connections = config.get("connections", [])
             for conn in connections:
                 if conn not in self.dependents:
                     self.dependents[conn] = set()
                 self.dependents[conn].add(agent_id)
-    
+
     def route(self, network, call_count, last_result, agent_stack):
         """Determine next agent to execute."""
         state = network.state
-        
+
         # First call - start with entry point or roots
         if call_count == 0:
             if self.entry_point:
@@ -333,28 +403,30 @@ class SwarmRouter(DeterministicRouter):
                 if roots:
                     state.current_agent = roots[0]
                     return self._get_agent_class(roots[0], agent_stack)
-        
+
         # Find next agent to execute
         for agent_id in self.agents_config:
             if agent_id in state.completed_agents:
                 continue
-            
+
             # Check if all dependencies are met
             deps = self.dependencies.get(agent_id, set())
             if deps.issubset(state.completed_agents):
                 state.current_agent = agent_id
                 return self._get_agent_class(agent_id, agent_stack)
-        
+
         # No more agents to execute
         return None
-    
-    def _get_agent_class(self, agent_id: str, agent_stack: List[type[Agent]]) -> type[Agent]:
+
+    def _get_agent_class(
+        self, agent_id: str, agent_stack: List[type[Agent]]
+    ) -> type[Agent]:
         """Get agent class for given agent ID."""
         # Find matching agent by name
         for agent_class in agent_stack:
             if hasattr(agent_class, "name") and agent_class.name == agent_id:
                 return agent_class
-        
+
         # Not found - this shouldn't happen
         return None
 
@@ -362,13 +434,13 @@ class SwarmRouter(DeterministicRouter):
 @final
 class SwarmTool(BaseTool):
     """Tool for executing agent networks using hanzo-agents SDK."""
-    
+
     @property
     @override
     def name(self) -> str:
         """Get the tool name."""
         return "swarm"
-    
+
     @property
     @override
     def description(self) -> str:
@@ -406,7 +478,7 @@ Models can be specified as:
 - CLI tools: 'claude_cli', 'codex_cli', 'gemini_cli', 'grok_cli'
 - Model URIs: 'model://anthropic/claude-3-opus'
 """
-    
+
     def __init__(
         self,
         permission_manager: PermissionManager,
@@ -421,13 +493,18 @@ Models can be specified as:
         self.permission_manager = permission_manager
         # Default to latest Claude Sonnet if no model specified
         from hanzo_mcp.tools.agent.code_auth import get_latest_claude_model
+
         self.model = model or f"anthropic/{get_latest_claude_model()}"
-        self.api_key = api_key or os.environ.get("ANTHROPIC_API_KEY") or os.environ.get("CLAUDE_API_KEY")
+        self.api_key = (
+            api_key
+            or os.environ.get("ANTHROPIC_API_KEY")
+            or os.environ.get("CLAUDE_API_KEY")
+        )
         self.base_url = base_url
         self.max_tokens = max_tokens
         self.agent_max_iterations = agent_max_iterations
         self.agent_max_tool_uses = agent_max_tool_uses
-        
+
         # Set up available tools for agents
         self.available_tools: list[BaseTool] = []
         self.available_tools.extend(
@@ -436,16 +513,16 @@ Models can be specified as:
         self.available_tools.extend(
             get_read_only_jupyter_tools(self.permission_manager)
         )
-        
+
         # Add edit tools
         self.available_tools.append(Edit(self.permission_manager))
         self.available_tools.append(MultiEdit(self.permission_manager))
-        
+
         # Add batch tool
         self.available_tools.append(
             BatchTool({t.name: t for t in self.available_tools})
         )
-    
+
     @override
     async def call(
         self,
@@ -455,7 +532,7 @@ Models can be specified as:
         """Execute the swarm tool."""
         tool_ctx = create_tool_context(ctx)
         await tool_ctx.set_tool_info(self.name)
-        
+
         # Extract parameters
         config = params.get("config", {})
         initial_query = params.get("query", "")
@@ -463,70 +540,76 @@ Models can be specified as:
         max_concurrent = params.get("max_concurrent", 10)
         use_memory = params.get("use_memory", False)
         memory_backend = params.get("memory_backend", "sqlite")
-        
+
         agents_config = config.get("agents", {})
-        
+
         if not agents_config:
             await tool_ctx.error("No agents provided")
             return "Error: At least one agent must be provided."
-        
+
         # hanzo-agents SDK is required (already imported above)
-        
-        await tool_ctx.info(f"Starting swarm execution with {len(agents_config)} agents using hanzo-agents SDK")
-        
-        # Create state
-        state = SwarmState(
-            config=config,
-            initial_query=initial_query,
-            context=context
+
+        await tool_ctx.info(
+            f"Starting swarm execution with {len(agents_config)} agents using hanzo-agents SDK"
         )
-        
+
+        # Create state
+        state = SwarmState(config=config, initial_query=initial_query, context=context)
+
         # Create agent classes dynamically
         agent_classes = []
         for agent_id, agent_config in agents_config.items():
             # Check for CLI agents
             model = agent_config.get("model", self.model)
-            
+
             cli_agents = {
                 "claude_cli": ClaudeCodeAgent,
                 "codex_cli": OpenAICodexAgent,
                 "gemini_cli": GeminiAgent,
                 "grok_cli": GrokAgent,
             }
-            
+
             if model in cli_agents:
                 # Use CLI agent
-                agent_class = type(f"Swarm{agent_id}", (cli_agents[model],), {
-                    "name": agent_id,
-                    "description": agent_config.get("role", f"Agent {agent_id}"),
-                    "agent_config": agent_config
-                })
+                agent_class = type(
+                    f"Swarm{agent_id}",
+                    (cli_agents[model],),
+                    {
+                        "name": agent_id,
+                        "description": agent_config.get("role", f"Agent {agent_id}"),
+                        "agent_config": agent_config,
+                    },
+                )
             else:
                 # Create dynamic SwarmAgent class
-                agent_class = type(f"Swarm{agent_id}", (SwarmAgent,), {
-                    "name": agent_id,
-                    "__init__": lambda self, aid=agent_id, acfg=agent_config: SwarmAgent.__init__(
-                        self,
-                        agent_id=aid,
-                        agent_config=acfg,
-                        available_tools=self.available_tools,
-                        permission_manager=self.permission_manager,
-                        ctx=ctx
-                    )
-                })
-            
+                agent_class = type(
+                    f"Swarm{agent_id}",
+                    (SwarmAgent,),
+                    {
+                        "name": agent_id,
+                        "__init__": lambda self, aid=agent_id, acfg=agent_config: SwarmAgent.__init__(
+                            self,
+                            agent_id=aid,
+                            agent_config=acfg,
+                            available_tools=self.available_tools,
+                            permission_manager=self.permission_manager,
+                            ctx=ctx,
+                        ),
+                    },
+                )
+
             agent_classes.append(agent_class)
-        
+
         # Create memory if requested
         memory_kv = None
         memory_vector = None
         if use_memory:
             memory_kv = create_memory_kv(memory_backend)
             memory_vector = create_memory_vector("simple")
-        
+
         # Create router
         router = SwarmRouter(config)
-        
+
         # Create network
         network = Network(
             state=state,
@@ -536,75 +619,77 @@ Models can be specified as:
             memory_vector=memory_vector,
             max_steps=self.agent_max_iterations * len(agents_config),
         )
-        
+
         # Execute
         try:
             final_state = await network.run()
-            
+
             # Format results
             return self._format_network_results(
                 agents_config,
                 final_state.agent_results,
                 final_state.execution_order,
-                config.get("entry_point")
+                config.get("entry_point"),
             )
-            
+
         except Exception as e:
             await tool_ctx.error(f"Swarm execution failed: {str(e)}")
             return f"Error: {str(e)}"
-    
+
     def _format_network_results(
         self,
         agents_config: Dict[str, Any],
         results: Dict[str, str],
         execution_order: List[str],
-        entry_point: Optional[str]
+        entry_point: Optional[str],
     ) -> str:
         """Format results from agent network execution."""
         output = ["Agent Network Execution Results (hanzo-agents SDK)"]
         output.append("=" * 80)
         output.append(f"Total agents: {len(agents_config)}")
         output.append(f"Completed: {len(results)}")
-        output.append(f"Failed: {len([r for r in results.values() if r.startswith('Error:')])}")
-        
+        output.append(
+            f"Failed: {len([r for r in results.values() if r.startswith('Error:')])}"
+        )
+
         if entry_point:
             output.append(f"Entry point: {entry_point}")
-        
+
         output.append(f"\nExecution Order: {' → '.join(execution_order)}")
         output.append("-" * 40)
-        
+
         # Detailed results
         output.append("\n\nDetailed Results:")
         output.append("=" * 80)
-        
+
         for agent_id in execution_order:
             if agent_id in results:
                 config = agents_config.get(agent_id, {})
                 role = config.get("role", "Agent")
                 model = config.get("model", "default")
-                
+
                 output.append(f"\n### {agent_id} ({role}) [{model}]")
                 output.append("-" * 40)
-                
+
                 result = results[agent_id]
                 if result.startswith("Error:"):
                     output.append(result)
                 else:
                     # Show first part of result
-                    lines = result.split('\n')
+                    lines = result.split("\n")
                     preview_lines = lines[:10]
                     output.extend(preview_lines)
-                    
+
                     if len(lines) > 10:
                         output.append(f"... ({len(lines) - 10} more lines)")
-        
+
         return "\n".join(output)
-    
+
     @override
     def register(self, mcp_server: FastMCP) -> None:
         """Register this swarm tool with the MCP server."""
         tool_self = self
-        
+
         @mcp_server.tool(name=self.name, description=self.description)
         async def swarm(
             ctx: MCPContext,
@@ -613,15 +698,15 @@ Models can be specified as:
             context: Optional[str] = None,
             max_concurrent: int = 10,
             use_memory: bool = False,
-            memory_backend: str = "sqlite"
+            memory_backend: str = "sqlite",
         ) -> str:
             # Convert to typed format
             typed_config = SwarmConfig(
                 agents=config.get("agents", {}),
                 entry_point=config.get("entry_point"),
-                topology=config.get("topology")
+                topology=config.get("topology"),
             )
-            
+
             return await tool_self.call(
                 ctx,
                 config=typed_config,
@@ -629,5 +714,5 @@ Models can be specified as:
                 context=context,
                 max_concurrent=max_concurrent,
                 use_memory=use_memory,
-                memory_backend=memory_backend
-            )
+                memory_backend=memory_backend,
+            )