hanzo-mcp 0.8.8__py3-none-any.whl → 0.9.0__py3-none-any.whl

hanzo_mcp/core/base_agent.py

@@ -28,11 +28,11 @@ 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
@@ -41,7 +41,7 @@ class AgentContext(Protocol):
 @dataclass
 class AgentConfig:
     """Configuration for agent execution."""
-    
+
     model: str = "claude-3-5-sonnet-20241022"
     timeout: int = 300
     max_retries: int = 3
@@ -49,7 +49,7 @@ class AgentConfig:
     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)
@@ -60,13 +60,13 @@ class AgentConfig:
 @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)."""
@@ -75,33 +75,33 @@ class AgentResult:
 
 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,
@@ -109,76 +109,76 @@ class BaseAgent(ABC, Generic[TContext]):
         **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,
@@ -187,47 +187,47 @@ class BaseAgent(ABC, Generic[TContext]):
         **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)
-        
+                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,
@@ -237,13 +237,13 @@ class BaseAgent(ABC, Generic[TContext]):
         **kwargs: Any,
     ) -> str:
         """Implementation-specific execution.
-        
+
         Args:
             prompt: The prompt
             context: Execution context
             env: Environment variables
             **kwargs: Additional parameters
-            
+
         Returns:
             Execution output
         """
@@ -252,35 +252,35 @@ class BaseAgent(ABC, Generic[TContext]):
 
 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,
@@ -289,21 +289,21 @@ class CLIAgent(BaseAgent[TContext]):
         **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,
@@ -313,7 +313,7 @@ class CLIAgent(BaseAgent[TContext]):
             cwd=str(self.config.working_dir) if self.config.working_dir else None,
             env=env,
         )
-        
+
         # Handle timeout
         try:
             stdout, stderr = await asyncio.wait_for(
@@ -323,18 +323,18 @@ class CLIAgent(BaseAgent[TContext]):
         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,
@@ -343,13 +343,13 @@ class APIAgent(BaseAgent[TContext]):
         **kwargs: Any,
     ) -> str:
         """Execute via API.
-        
+
         Args:
             prompt: The prompt
             context: Execution context
             env: Environment variables
             **kwargs: Additional parameters
-            
+
         Returns:
             API response
         """
@@ -360,36 +360,36 @@ class APIAgent(BaseAgent[TContext]):
 
 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,
@@ -398,13 +398,13 @@ class AgentOrchestrator:
         **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
         """
@@ -414,25 +414,25 @@ class AgentOrchestrator:
                 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(
@@ -441,12 +441,12 @@ class AgentOrchestrator:
                     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,
@@ -454,58 +454,58 @@ class AgentOrchestrator:
         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
 
 
@@ -517,4 +517,4 @@ __all__ = [
     "CLIAgent",
     "APIAgent",
     "AgentOrchestrator",
-]
+]