diff --git a/src/agents/coordination.py b/src/agents/coordination.py new file mode 100644 index 0000000..b582402 --- /dev/null +++ b/src/agents/coordination.py @@ -0,0 +1,407 @@ +""" +Multi-agent coordination engine. + +Orchestrates delegation from Tatlock to expert agents (Librarian, etc.) +based on Steward recommendations. Handles: +- Routing tasks to appropriate agents +- Parallel and sequential execution +- Result aggregation +- Error handling and graceful degradation +""" +import asyncio +import time +from typing import Any, AsyncGenerator, Optional + +from src.agents.librarian import run_librarian, run_librarian_stream +from src.agents.protocol import ( + AgentError, + AgentRequest, + AgentResponse, + AgentTimeoutError, + AgentUnavailableError, + CoordinationResult, + DelegationIntent, + DelegationReason, + ToolCallRecord, +) +from src.core.household_registry import get_household_registry +from src.core.logging_config import get_logger + +logger = get_logger(__name__) + + +# Agent execution functions registry +AGENT_EXECUTORS: dict[str, Any] = { + "librarian": run_librarian, +} + +AGENT_STREAM_EXECUTORS: dict[str, Any] = { + "librarian": run_librarian_stream, +} + + +class CoordinationEngine: + """ + Coordinates multi-agent task execution. + + Routes tasks from Tatlock to appropriate expert agents, + handles execution, and aggregates results. + """ + + def __init__(self): + """Initialize the coordination engine.""" + self.registry = get_household_registry() + logger.info("coordination_engine_initialized") + + def get_available_agents(self) -> list[str]: + """ + Get list of available expert agents. + + Returns: + List of agent names that can accept delegations + """ + available = [] + for name in self.registry.list_members(): + member = self.registry.get_member(name) + if member and member.agent is not None: + available.append(name) + return available + + def can_delegate_to(self, agent_name: str) -> bool: + """ + Check if delegation to an agent is possible. + + Args: + agent_name: Name of the target agent + + Returns: + True if agent is available and can accept tasks + """ + if agent_name not in AGENT_EXECUTORS: + return False + + member = self.registry.get_member(agent_name) + return member is not None and member.agent is not None + + async def execute_delegation( + self, + intent: DelegationIntent, + context: str = "", + message_history: Optional[list[Any]] = None, + ) -> AgentResponse: + """ + Execute a single delegation to an expert agent. + + Args: + intent: The delegation intent with task details + context: Additional context for the agent + message_history: Optional conversation history + + Returns: + AgentResponse with results + + Raises: + AgentUnavailableError: If agent is not available + AgentTimeoutError: If execution times out + AgentError: For other execution errors + """ + start_time = time.time() + agent_name = intent.target_agent + + logger.info( + "delegation_started", + agent=agent_name, + task=intent.task[:100], + reason=intent.reason.value, + ) + + # Check if agent is available + if not self.can_delegate_to(agent_name): + raise AgentUnavailableError( + f"Agent '{agent_name}' is not available for delegation", + agent_name=agent_name, + ) + + # Get the executor + executor = AGENT_EXECUTORS.get(agent_name) + if not executor: + raise AgentUnavailableError( + f"No executor found for agent '{agent_name}'", + agent_name=agent_name, + ) + + try: + # Build the request + request = AgentRequest( + task=intent.task, + context=context, + delegation_reason=intent.reason, + ) + + # Execute with timeout + timeout = request.timeout_seconds or 60 + + result = await asyncio.wait_for( + executor( + task=request.task, + context=request.context, + message_history=message_history, + ), + timeout=timeout, + ) + + duration_ms = int((time.time() - start_time) * 1000) + + logger.info( + "delegation_completed", + agent=agent_name, + duration_ms=duration_ms, + output_length=len(result), + ) + + return AgentResponse( + success=True, + result=result, + reasoning=f"Delegated to {agent_name}: {intent.expected_outcome}", + duration_ms=duration_ms, + ) + + except asyncio.TimeoutError: + duration_ms = int((time.time() - start_time) * 1000) + logger.error( + "delegation_timeout", + agent=agent_name, + duration_ms=duration_ms, + ) + raise AgentTimeoutError( + f"Agent '{agent_name}' timed out after {duration_ms}ms", + agent_name=agent_name, + ) + + except Exception as e: + duration_ms = int((time.time() - start_time) * 1000) + logger.error( + "delegation_error", + agent=agent_name, + error=str(e), + duration_ms=duration_ms, + exc_info=True, + ) + return AgentResponse( + success=False, + result="", + error_message=str(e), + duration_ms=duration_ms, + ) + + async def execute_delegation_stream( + self, + intent: DelegationIntent, + context: str = "", + message_history: Optional[list[Any]] = None, + ) -> AsyncGenerator[str, None]: + """ + Execute a delegation with streaming output. + + Args: + intent: The delegation intent with task details + context: Additional context for the agent + message_history: Optional conversation history + + Yields: + Text deltas from the agent + + Raises: + AgentUnavailableError: If agent is not available + """ + agent_name = intent.target_agent + + logger.info( + "delegation_stream_started", + agent=agent_name, + task=intent.task[:100], + ) + + # Check if agent is available + if agent_name not in AGENT_STREAM_EXECUTORS: + raise AgentUnavailableError( + f"Agent '{agent_name}' does not support streaming", + agent_name=agent_name, + ) + + executor = AGENT_STREAM_EXECUTORS[agent_name] + + try: + async for delta in executor( + task=intent.task, + context=context, + message_history=message_history, + ): + yield delta + + logger.info("delegation_stream_completed", agent=agent_name) + + except Exception as e: + logger.error( + "delegation_stream_error", + agent=agent_name, + error=str(e), + exc_info=True, + ) + yield f"\n\n[Error from {agent_name}: {str(e)}]" + + async def coordinate( + self, + intents: list[DelegationIntent], + context: str = "", + message_history: Optional[list[Any]] = None, + ) -> CoordinationResult: + """ + Coordinate execution of multiple delegations. + + Handles parallel execution for independent tasks and + sequential execution for dependent tasks. + + Args: + intents: List of delegation intents to execute + context: Shared context for all agents + message_history: Optional conversation history + + Returns: + CoordinationResult with aggregated results + """ + start_time = time.time() + agent_responses: dict[str, AgentResponse] = {} + agents_consulted: list[str] = [] + + logger.info( + "coordination_started", + intent_count=len(intents), + agents=[i.target_agent for i in intents], + ) + + # Sort by priority + sorted_intents = sorted(intents, key=lambda x: x.priority) + + # Group by dependencies (simple version: sequential for now) + # TODO: Implement parallel execution for independent tasks + for intent in sorted_intents: + try: + response = await self.execute_delegation( + intent=intent, + context=context, + message_history=message_history, + ) + agent_responses[intent.target_agent] = response + if response.success: + agents_consulted.append(intent.target_agent) + + except AgentError as e: + agent_responses[intent.target_agent] = AgentResponse( + success=False, + result="", + error_message=str(e), + ) + + # Aggregate results + successful_results = [ + r.result for r in agent_responses.values() if r.success and r.result + ] + + final_response = "\n\n---\n\n".join(successful_results) if successful_results else "" + + total_duration = int((time.time() - start_time) * 1000) + + logger.info( + "coordination_completed", + total_duration_ms=total_duration, + agents_consulted=agents_consulted, + success_count=len(successful_results), + ) + + return CoordinationResult( + final_response=final_response, + agent_responses=agent_responses, + delegation_intents=intents, + total_duration_ms=total_duration, + agents_consulted=agents_consulted, + ) + + +# Global coordination engine instance +_coordination_engine: Optional[CoordinationEngine] = None + + +def get_coordination_engine() -> CoordinationEngine: + """Get the global coordination engine instance.""" + global _coordination_engine + if _coordination_engine is None: + _coordination_engine = CoordinationEngine() + return _coordination_engine + + +async def delegate_to_librarian( + task: str, + context: str = "", + reason: DelegationReason = DelegationReason.DOMAIN_EXPERTISE, + message_history: Optional[list[Any]] = None, +) -> AgentResponse: + """ + Convenience function to delegate a task to The Librarian. + + Args: + task: Research task description + context: Additional context + reason: Why delegating to Librarian + message_history: Optional conversation history + + Returns: + AgentResponse with research results + """ + engine = get_coordination_engine() + + intent = DelegationIntent( + target_agent="librarian", + task=task, + reason=reason, + expected_outcome="Research findings and relevant information", + ) + + return await engine.execute_delegation( + intent=intent, + context=context, + message_history=message_history, + ) + + +async def delegate_to_librarian_stream( + task: str, + context: str = "", + message_history: Optional[list[Any]] = None, +) -> AsyncGenerator[str, None]: + """ + Convenience function to delegate to Librarian with streaming. + + Args: + task: Research task description + context: Additional context + message_history: Optional conversation history + + Yields: + Text deltas from The Librarian + """ + engine = get_coordination_engine() + + intent = DelegationIntent( + target_agent="librarian", + task=task, + reason=DelegationReason.DOMAIN_EXPERTISE, + expected_outcome="Research findings", + ) + + async for delta in engine.execute_delegation_stream( + intent=intent, + context=context, + message_history=message_history, + ): + yield delta