python-sdk/pkg/hanzo-tools-agent/hanzo_tools/agent/swarm_tool.py

"""Swarm tool implementation using hanzo-agents SDK.

This module implements the SwarmTool that leverages the hanzo-agents SDK
for sophisticated multi-agent orchestration with flexible network topologies.
"""

import os
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 hanzo_tools.core import auto_timeout

# Import hanzo-agents SDK with fallback
try:
    from hanzo_agents import (
        Tool,
        Agent,
        State,
        Router,
        History,
        Network,
        ToolCall,
        ModelRegistry,
        InferenceResult,
    )

    HANZO_AGENTS_AVAILABLE = True
except ImportError:
    # hanzo-agents not installed - types raise ImportError on use
    HANZO_AGENTS_AVAILABLE = False

    class _RequiresHanzoAgents(type):
        """Raises ImportError on instantiation."""

        def __call__(cls, *args, **kwargs):
            raise ImportError(f"{cls.__name__} requires: pip install hanzo-agents")

    class Agent(metaclass=_RequiresHanzoAgents):
        pass

    class State(metaclass=_RequiresHanzoAgents):
        pass

    class Network(metaclass=_RequiresHanzoAgents):
        pass

    class Tool(metaclass=_RequiresHanzoAgents):
        pass

    class History(metaclass=_RequiresHanzoAgents):
        pass

    class ModelRegistry(metaclass=_RequiresHanzoAgents):
        pass

    class InferenceResult(metaclass=_RequiresHanzoAgents):
        pass

    class ToolCall(metaclass=_RequiresHanzoAgents):
        pass

    class Router(metaclass=_RequiresHanzoAgents):
        pass


# Import optional components with fallbacks
try:
    from hanzo_agents import LLMRouter, HybridRouter, DeterministicRouter
except ImportError:
    try:
        # Try core module import
        from hanzo_agents.core.router import (
            LLMRouter,
            HybridRouter,
            DeterministicRouter,
        )
    except ImportError:

        class DeterministicRouter(metaclass=_RequiresHanzoAgents):
            pass

        class LLMRouter(metaclass=_RequiresHanzoAgents):
            pass

        class HybridRouter(metaclass=_RequiresHanzoAgents):
            pass


def _requires_hanzo_agents(name: str):
    """Create function that raises ImportError."""

    def fn(*args, **kwargs):
        raise ImportError(f"{name} requires: pip install hanzo-agents")

    fn.__name__ = name
    return fn


try:
    from hanzo_agents import create_memory_kv, create_memory_vector
except ImportError:
    try:
        from hanzo_agents.core.memory import create_memory_kv, create_memory_vector
    except ImportError:
        create_memory_kv = _requires_hanzo_agents("create_memory_kv")
        create_memory_vector = _requires_hanzo_agents("create_memory_vector")


try:
    from hanzo_agents import sequential_router, conditional_router, state_based_router
except ImportError:
    try:
        from hanzo_agents.core.router import (
            sequential_router,
            conditional_router,
            state_based_router,
        )
    except ImportError:
        sequential_router = _requires_hanzo_agents("sequential_router")
        conditional_router = _requires_hanzo_agents("conditional_router")
        state_based_router = _requires_hanzo_agents("state_based_router")


try:
    from hanzo_agents.core.cli_agent import (
        GrokAgent,
        GeminiAgent,
        ClaudeCodeAgent,
        OpenAICodexAgent,
    )
except ImportError:

    class ClaudeCodeAgent(metaclass=_RequiresHanzoAgents):
        pass

    class OpenAICodexAgent(metaclass=_RequiresHanzoAgents):
        pass

    class GeminiAgent(metaclass=_RequiresHanzoAgents):
        pass

    class GrokAgent(metaclass=_RequiresHanzoAgents):
        pass


from hanzo_tools.fs import EditTool, get_read_only_filesystem_tools
from hanzo_tools.core import BaseTool, PermissionManager, create_tool_context
from hanzo_tools.jupyter import get_read_only_jupyter_tools

from .agent_tool import MCPAgent


class AgentNode(TypedDict):
    """Node in the agent network."""

    id: str
    query: str
    model: Optional[str]
    role: Optional[str]
    connections: Optional[List[str]]
    receives_from: Optional[List[str]]
    file_path: Optional[str]


class SwarmConfig(TypedDict):
    """Configuration for an agent network."""

    agents: Dict[str, AgentNode]
    entry_point: Optional[str]
    topology: Optional[str]


class SwarmToolParams(TypedDict):
    """Parameters for the SwarmTool."""

    config: SwarmConfig
    query: str
    context: Optional[str]
    max_concurrent: Optional[int]
    use_memory: Optional[bool]
    memory_backend: Optional[str]


class SwarmState(State):
    """State for swarm execution."""

    def __init__(
        self, config: SwarmConfig, initial_query: str, context: Optional[str] = None
    ):
        """Initialize swarm state."""
        super().__init__()
        self.config = config
        self.initial_query = initial_query
        self.context = context
        self.agent_results = {}
        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,
            }
        )
        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"),
        )
        state.agent_results = data.get("agent_results", {})
        state.completed_agents = set(data.get("completed_agents", []))
        state.current_agent = data.get("current_agent")
        state.execution_order = data.get("execution_order", [])
        return 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,
    ):
        """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,
        )

    def _normalize_model(self, model: str) -> str:
        """Normalize model names to full format."""
        model_map = {
            "claude-3-5-sonnet": "model://anthropic/claude-3-5-sonnet-20241022",
            "claude-3-opus": "model://anthropic/claude-3-opus-20240229",
            "gpt-4o": "model://openai/gpt-4o",
            "gpt-4": "model://openai/gpt-4",
            "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:
        """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:
            inputs = {}
            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},
        ]

        # 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,
            content=response,
            metadata={
                "agent_id": self.name,
                "role": self.agent_config.get("role"),
                "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:
                state.current_agent = self.entry_point
                return self._get_agent_class(self.entry_point, agent_stack)
            else:
                # Find roots (no dependencies)
                roots = [aid for aid, deps in self.dependencies.items() if not deps]
                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]:
        """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


@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:
        """Get the tool description."""
        return """Execute a network of AI agents with flexible connection topologies.

This tool enables sophisticated agent orchestration where agents can be connected
in various network patterns. Each agent can pass results to connected agents,
enabling complex workflows.

Features:
- Flexible agent networks (tree, DAG, pipeline, star, mesh)
- Each agent can use different models (Claude, GPT-4, Gemini, etc.)
- Agents automatically pass results to connected agents
- Parallel execution with dependency management
- Full editing capabilities for each agent
- Memory and state management via hanzo-agents SDK

Common Topologies:
1. Tree (Architect pattern):
   architect → [frontend, backend, database] → reviewer

2. Pipeline (Sequential processing):
   analyzer → planner → implementer → tester → reviewer

3. Star (Central coordinator):
   coordinator ← → [agent1, agent2, agent3, agent4]

4. DAG (Complex dependencies):
   Multiple agents with custom connections

Models can be specified as:
- Full: 'anthropic/claude-3-5-sonnet-20241022'
- Short: 'claude-3-5-sonnet', 'gpt-4o', 'gemini-1.5-pro'
- CLI tools: 'claude_cli', 'codex_cli', 'gemini_cli', 'grok_cli'
- Model URIs: 'model://anthropic/claude-3-opus'
"""

    def __init__(
        self,
        permission_manager: PermissionManager,
        model: str | None = None,
        api_key: str | None = None,
        base_url: str | None = None,
        max_tokens: int | None = None,
        agent_max_iterations: int = 10,
        agent_max_tool_uses: int = 30,
    ):
        """Initialize the swarm tool."""
        self.permission_manager = permission_manager
        # Default to latest Claude Sonnet if no model specified
        from .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.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(
            get_read_only_filesystem_tools(self.permission_manager)
        )
        self.available_tools.extend(
            get_read_only_jupyter_tools(self.permission_manager)
        )

        # Add edit tool
        self.available_tools.append(EditTool(self.permission_manager))

    @override
    @auto_timeout("swarm")
    async def call(
        self,
        ctx: MCPContext,
        **params: Unpack[SwarmToolParams],
    ) -> str:
        """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", "")
        context = params.get("context", "")
        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)

        # 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,
                    },
                )
            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_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,
            agents=agent_classes,
            router=router,
            memory_kv=memory_kv,
            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"),
            )

        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],
    ) -> 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:')])}"
        )

        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")
                    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,
            config: dict[str, Any],
            query: str,
            context: Optional[str] = None,
            max_concurrent: int = 10,
            use_memory: bool = False,
            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"),
            )

            return await tool_self.call(
                ctx,
                config=typed_config,
                query=query,
                context=context,
                max_concurrent=max_concurrent,
                use_memory=use_memory,
                memory_backend=memory_backend,
            )