hanzo-mcp 0.8.2__py3-none-any.whl → 0.8.4__py3-none-any.whl

hanzo_mcp/core/base_agent.py

@@ -0,0 +1,520 @@
+"""Base Agent - Unified foundation for all AI agent implementations.
+
+This module provides the single base class for all agent operations,
+following DRY principles and ensuring consistent behavior across all agents.
+"""
+
+from __future__ import annotations
+
+import os
+import asyncio
+import logging
+from abc import ABC, abstractmethod
+from typing import Any, Dict, List, Generic, TypeVar, Optional, Protocol, runtime_checkable
+from pathlib import Path
+from datetime import datetime
+from dataclasses import field, dataclass
+
+from .model_registry import ModelConfig, registry
+
+logger = logging.getLogger(__name__)
+
+
+# Type variables for generic context
+TContext = TypeVar("TContext")
+TResult = TypeVar("TResult")
+
+
+@runtime_checkable
+class AgentContext(Protocol):
+    """Protocol for agent execution context."""
+    
+    async def log(self, message: str, level: str = "info") -> None:
+        """Log a message."""
+        pass
+    
+    async def progress(self, message: str, percentage: Optional[float] = None) -> None:
+        """Report progress."""
+        pass
+
+
+@dataclass
+class AgentConfig:
+    """Configuration for agent execution."""
+    
+    model: str = "claude-3-5-sonnet-20241022"
+    timeout: int = 300
+    max_retries: int = 3
+    working_dir: Optional[Path] = None
+    environment: Dict[str, str] = field(default_factory=dict)
+    stream_output: bool = False
+    use_worktree: bool = False
+    
+    def __post_init__(self) -> None:
+        """Resolve model name and validate configuration."""
+        self.model = registry.resolve(self.model)
+        if self.working_dir and not isinstance(self.working_dir, Path):
+            self.working_dir = Path(self.working_dir)
+
+
+@dataclass
+class AgentResult:
+    """Result from agent execution."""
+    
+    success: bool
+    output: Optional[str] = None
+    error: Optional[str] = None
+    duration: Optional[float] = None
+    metadata: Dict[str, Any] = field(default_factory=dict)
+    
+    @property
+    def content(self) -> str:
+        """Get the primary content (output or error)."""
+        return self.output if self.success else (self.error or "Unknown error")
+
+
+class BaseAgent(ABC, Generic[TContext]):
+    """Base class for all AI agents.
+    
+    This is the single foundation for all agent implementations,
+    ensuring consistent behavior and eliminating code duplication.
+    """
+    
+    def __init__(self, config: Optional[AgentConfig] = None) -> None:
+        """Initialize agent with configuration.
+        
+        Args:
+            config: Agent configuration
+        """
+        self.config = config or AgentConfig()
+        self._start_time: Optional[datetime] = None
+        self._end_time: Optional[datetime] = None
+    
+    @property
+    @abstractmethod
+    def name(self) -> str:
+        """Agent name."""
+        pass
+    
+    @property
+    @abstractmethod
+    def description(self) -> str:
+        """Agent description."""
+        pass
+    
+    async def execute(
+        self,
+        prompt: str,
+        context: Optional[TContext] = None,
+        **kwargs: Any,
+    ) -> AgentResult:
+        """Execute agent with prompt.
+        
+        Args:
+            prompt: The prompt or task
+            context: Execution context
+            **kwargs: Additional parameters
+            
+        Returns:
+            Agent execution result
+        """
+        self._start_time = datetime.now()
+        
+        try:
+            # Setup environment
+            env = self._prepare_environment()
+            
+            # Log start
+            if context and isinstance(context, AgentContext):
+                await context.log(f"Starting {self.name} with model {self.config.model}")
+            
+            # Execute with retries
+            result = await self._execute_with_retries(prompt, context, env, **kwargs)
+            
+            # Calculate duration
+            self._end_time = datetime.now()
+            duration = (self._end_time - self._start_time).total_seconds()
+            
+            return AgentResult(
+                success=True,
+                output=result,
+                duration=duration,
+                metadata={"model": self.config.model, "agent": self.name},
+            )
+            
+        except Exception as e:
+            self._end_time = datetime.now()
+            duration = (self._end_time - self._start_time).total_seconds() if self._start_time else None
+            
+            logger.error(f"Agent {self.name} failed: {e}")
+            
+            return AgentResult(
+                success=False,
+                error=str(e),
+                duration=duration,
+                metadata={"model": self.config.model, "agent": self.name},
+            )
+    
+    def _prepare_environment(self) -> Dict[str, str]:
+        """Prepare environment variables for execution.
+        
+        Returns:
+            Environment variables dictionary
+        """
+        env = os.environ.copy()
+        
+        # Add model-specific API key
+        model_config = registry.get(self.config.model)
+        if model_config and model_config.api_key_env:
+            key_var = model_config.api_key_env
+            if key_var in os.environ:
+                env[key_var] = os.environ[key_var]
+        
+        # Add Hanzo unified auth
+        if "HANZO_API_KEY" in os.environ:
+            env["HANZO_API_KEY"] = os.environ["HANZO_API_KEY"]
+        
+        # Add custom environment
+        env.update(self.config.environment)
+        
+        return env
+    
+    async def _execute_with_retries(
+        self,
+        prompt: str,
+        context: Optional[TContext],
+        env: Dict[str, str],
+        **kwargs: Any,
+    ) -> str:
+        """Execute with retry logic.
+        
+        Args:
+            prompt: The prompt
+            context: Execution context
+            env: Environment variables
+            **kwargs: Additional parameters
+            
+        Returns:
+            Execution output
+            
+        Raises:
+            Exception: If all retries fail
+        """
+        last_error = None
+        
+        for attempt in range(self.config.max_retries):
+            try:
+                # Call the implementation
+                result = await self._execute_impl(prompt, context, env, **kwargs)
+                return result
+                
+            except asyncio.TimeoutError:
+                last_error = f"Timeout after {self.config.timeout} seconds"
+                if context and isinstance(context, AgentContext):
+                    await context.log(f"Attempt {attempt + 1} timed out", "warning")
+                    
+            except Exception as e:
+                last_error = str(e)
+                if context and isinstance(context, AgentContext):
+                    await context.log(f"Attempt {attempt + 1} failed: {e}", "warning")
+                
+                # Don't retry on certain errors
+                if "unauthorized" in str(e).lower() or "forbidden" in str(e).lower():
+                    raise
+            
+            # Wait before retry (exponential backoff)
+            if attempt < self.config.max_retries - 1:
+                await asyncio.sleep(2 ** attempt)
+        
+        raise Exception(f"All {self.config.max_retries} attempts failed. Last error: {last_error}")
+    
+    @abstractmethod
+    async def _execute_impl(
+        self,
+        prompt: str,
+        context: Optional[TContext],
+        env: Dict[str, str],
+        **kwargs: Any,
+    ) -> str:
+        """Implementation-specific execution.
+        
+        Args:
+            prompt: The prompt
+            context: Execution context
+            env: Environment variables
+            **kwargs: Additional parameters
+            
+        Returns:
+            Execution output
+        """
+        pass
+
+
+class CLIAgent(BaseAgent[TContext]):
+    """Base class for CLI-based agents."""
+    
+    @property
+    @abstractmethod
+    def cli_command(self) -> str:
+        """CLI command to execute."""
+        pass
+    
+    def build_command(self, prompt: str, **kwargs: Any) -> List[str]:
+        """Build the CLI command.
+        
+        Args:
+            prompt: The prompt
+            **kwargs: Additional parameters
+            
+        Returns:
+            Command arguments list
+        """
+        command = [self.cli_command]
+        
+        # Add model if specified
+        model_config = registry.get(self.config.model)
+        if model_config:
+            command.extend(["--model", model_config.full_name])
+        
+        # Add prompt
+        command.append(prompt)
+        
+        return command
+    
+    async def _execute_impl(
+        self,
+        prompt: str,
+        context: Optional[TContext],
+        env: Dict[str, str],
+        **kwargs: Any,
+    ) -> str:
+        """Execute CLI command.
+        
+        Args:
+            prompt: The prompt
+            context: Execution context
+            env: Environment variables
+            **kwargs: Additional parameters
+            
+        Returns:
+            Command output
+        """
+        command = self.build_command(prompt, **kwargs)
+        
+        # Determine if we need stdin
+        needs_stdin = self.cli_command in ["claude", "cline"]
+        
+        # Execute command
+        process = await asyncio.create_subprocess_exec(
+            *command,
+            stdin=asyncio.subprocess.PIPE if needs_stdin else None,
+            stdout=asyncio.subprocess.PIPE,
+            stderr=asyncio.subprocess.PIPE,
+            cwd=str(self.config.working_dir) if self.config.working_dir else None,
+            env=env,
+        )
+        
+        # Handle timeout
+        try:
+            stdout, stderr = await asyncio.wait_for(
+                process.communicate(prompt.encode() if needs_stdin else None),
+                timeout=self.config.timeout,
+            )
+        except asyncio.TimeoutError:
+            process.kill()
+            raise asyncio.TimeoutError(f"Command timed out after {self.config.timeout} seconds")
+        
+        # Check for errors
+        if process.returncode != 0:
+            error_msg = stderr.decode() if stderr else "Command failed"
+            raise Exception(error_msg)
+        
+        return stdout.decode()
+
+
+class APIAgent(BaseAgent[TContext]):
+    """Base class for API-based agents."""
+    
+    async def _execute_impl(
+        self,
+        prompt: str,
+        context: Optional[TContext],
+        env: Dict[str, str],
+        **kwargs: Any,
+    ) -> str:
+        """Execute via API.
+        
+        Args:
+            prompt: The prompt
+            context: Execution context
+            env: Environment variables
+            **kwargs: Additional parameters
+            
+        Returns:
+            API response
+        """
+        # This would be implemented by specific API agents
+        # using the appropriate client library
+        raise NotImplementedError("API agents must implement _execute_impl")
+
+
+class AgentOrchestrator:
+    """Orchestrator for managing multiple agents."""
+    
+    def __init__(self, default_config: Optional[AgentConfig] = None) -> None:
+        """Initialize orchestrator.
+        
+        Args:
+            default_config: Default configuration for agents
+        """
+        self.default_config = default_config or AgentConfig()
+        self._agents: Dict[str, BaseAgent] = {}
+        self._semaphore: Optional[asyncio.Semaphore] = None
+    
+    def register(self, agent: BaseAgent) -> None:
+        """Register an agent.
+        
+        Args:
+            agent: Agent to register
+        """
+        self._agents[agent.name] = agent
+    
+    def get_agent(self, name: str) -> Optional[BaseAgent]:
+        """Get agent by name.
+        
+        Args:
+            name: Agent name
+            
+        Returns:
+            Agent instance or None
+        """
+        return self._agents.get(name)
+    
+    async def execute_single(
+        self,
+        agent_name: str,
+        prompt: str,
+        context: Optional[Any] = None,
+        **kwargs: Any,
+    ) -> AgentResult:
+        """Execute single agent.
+        
+        Args:
+            agent_name: Name of agent to use
+            prompt: The prompt
+            context: Execution context
+            **kwargs: Additional parameters
+            
+        Returns:
+            Execution result
+        """
+        agent = self.get_agent(agent_name)
+        if not agent:
+            return AgentResult(
+                success=False,
+                error=f"Agent '{agent_name}' not found",
+            )
+        
+        return await agent.execute(prompt, context, **kwargs)
+    
+    async def execute_parallel(
+        self,
+        tasks: List[Dict[str, Any]],
+        max_concurrent: int = 5,
+    ) -> List[AgentResult]:
+        """Execute multiple agents in parallel.
+        
+        Args:
+            tasks: List of task definitions
+            max_concurrent: Maximum concurrent executions
+            
+        Returns:
+            List of results
+        """
+        self._semaphore = asyncio.Semaphore(max_concurrent)
+        
+        async def run_with_semaphore(task: Dict[str, Any]) -> AgentResult:
+            async with self._semaphore:
+                return await self.execute_single(
+                    task["agent"],
+                    task["prompt"],
+                    task.get("context"),
+                    **task.get("kwargs", {}),
+                )
+        
+        return await asyncio.gather(
+            *[run_with_semaphore(task) for task in tasks],
+            return_exceptions=False,
+        )
+    
+    async def execute_consensus(
+        self,
+        prompt: str,
+        agents: List[str],
+        threshold: float = 0.66,
+    ) -> Dict[str, Any]:
+        """Execute consensus operation with multiple agents.
+        
+        Args:
+            prompt: The prompt
+            agents: List of agent names
+            threshold: Agreement threshold
+            
+        Returns:
+            Consensus results
+        """
+        # Execute all agents in parallel
+        tasks = [{"agent": agent, "prompt": prompt} for agent in agents]
+        results = await self.execute_parallel(tasks)
+        
+        # Analyze consensus
+        successful = [r for r in results if r.success]
+        agreement = len(successful) / len(results) if results else 0
+        
+        return {
+            "consensus_reached": agreement >= threshold,
+            "agreement_score": agreement,
+            "individual_results": results,
+            "agents_used": agents,
+        }
+    
+    async def execute_chain(
+        self,
+        initial_prompt: str,
+        agents: List[str],
+    ) -> List[AgentResult]:
+        """Execute agents in a chain, passing output forward.
+        
+        Args:
+            initial_prompt: Initial prompt
+            agents: List of agent names
+            
+        Returns:
+            List of results from each step
+        """
+        results = []
+        current_prompt = initial_prompt
+        
+        for agent_name in agents:
+            result = await self.execute_single(agent_name, current_prompt)
+            results.append(result)
+            
+            if result.success and result.output:
+                # Use output as input for next agent
+                current_prompt = f"Review and improve:\n{result.output}"
+            else:
+                # Chain broken
+                break
+        
+        return results
+
+
+__all__ = [
+    "AgentContext",
+    "AgentConfig",
+    "AgentResult",
+    "BaseAgent",
+    "CLIAgent",
+    "APIAgent",
+    "AgentOrchestrator",
+]