""" Advanced Agent Orchestrator for Cidadão.AI. Manages complex agent coordination patterns and workflows. """ import asyncio import time from datetime import datetime, timedelta from typing import Any, Dict, List, Optional, Callable, Tuple from dataclasses import dataclass, field from enum import Enum from collections import defaultdict import weakref from src.core import get_logger from src.agents.deodoro import BaseAgent, AgentContext, AgentMessage, AgentResponse from src.services.agent_lazy_loader import agent_lazy_loader from src.services.cache_service import CacheService from src.core.exceptions import OrchestrationError logger = get_logger("agent.orchestrator") class OrchestrationPattern(Enum): """Orchestration patterns supported by the system.""" SEQUENTIAL = "sequential" PARALLEL = "parallel" FAN_OUT_FAN_IN = "fan_out_fan_in" CONDITIONAL = "conditional" SAGA = "saga" MAP_REDUCE = "map_reduce" EVENT_DRIVEN = "event_driven" class CircuitState(Enum): """Circuit breaker states.""" CLOSED = "closed" OPEN = "open" HALF_OPEN = "half_open" @dataclass class CircuitBreaker: """Circuit breaker for agent fault tolerance.""" failure_threshold: int = 5 recovery_timeout: int = 60 # seconds half_open_requests: int = 3 state: CircuitState = CircuitState.CLOSED failure_count: int = 0 last_failure_time: Optional[datetime] = None success_count: int = 0 @dataclass class WorkflowStep: """Represents a step in an orchestrated workflow.""" step_id: str agent_name: str action: str input_mapping: Dict[str, str] = field(default_factory=dict) output_mapping: Dict[str, str] = field(default_factory=dict) conditions: Dict[str, Any] = field(default_factory=dict) retry_config: Dict[str, Any] = field(default_factory=dict) timeout: int = 300 # seconds @dataclass class WorkflowDefinition: """Defines an orchestrated workflow.""" workflow_id: str name: str pattern: OrchestrationPattern steps: List[WorkflowStep] metadata: Dict[str, Any] = field(default_factory=dict) timeout: int = 1800 # 30 minutes @dataclass class OrchestrationMetrics: """Metrics for orchestration performance.""" total_executions: int = 0 successful_executions: int = 0 failed_executions: int = 0 total_duration_seconds: float = 0.0 agent_execution_times: Dict[str, List[float]] = field(default_factory=lambda: defaultdict(list)) pattern_usage: Dict[str, int] = field(default_factory=lambda: defaultdict(int)) class EventBus: """Simple event bus for event-driven choreography.""" def __init__(self): self._handlers: Dict[str, List[Callable]] = defaultdict(list) self._async_handlers: Dict[str, List[Callable]] = defaultdict(list) def on(self, event_name: str, handler: Callable): """Register an event handler.""" if asyncio.iscoroutinefunction(handler): self._async_handlers[event_name].append(handler) else: self._handlers[event_name].append(handler) async def emit(self, event_name: str, data: Any = None): """Emit an event to all registered handlers.""" event = {"name": event_name, "data": data, "timestamp": datetime.utcnow()} # Call sync handlers for handler in self._handlers.get(event_name, []): try: handler(event) except Exception as e: logger.error(f"Error in event handler: {e}") # Call async handlers tasks = [] for handler in self._async_handlers.get(event_name, []): tasks.append(handler(event)) if tasks: await asyncio.gather(*tasks, return_exceptions=True) class AgentOrchestrator: """Advanced orchestrator for multi-agent coordination.""" def __init__(self): self.logger = logger self._workflows: Dict[str, WorkflowDefinition] = {} self._circuit_breakers: Dict[str, CircuitBreaker] = {} self._metrics = OrchestrationMetrics() self._event_bus = EventBus() self._cache = CacheService() self._agent_capabilities: Dict[str, List[str]] = {} self._running_workflows: weakref.WeakValueDictionary = weakref.WeakValueDictionary() async def initialize(self): """Initialize the orchestrator.""" self.logger.info("Initializing Agent Orchestrator") # Discover agent capabilities await self._discover_agent_capabilities() # Register default workflows self._register_default_workflows() async def _discover_agent_capabilities(self): """Discover capabilities of all available agents.""" try: agents = await agent_lazy_loader.list_agents() for agent_info in agents: agent = await agent_lazy_loader.get_agent(agent_info["name"]) if hasattr(agent, 'capabilities'): self._agent_capabilities[agent_info["name"]] = agent.capabilities except Exception as e: self.logger.error(f"Error discovering agent capabilities: {e}") def _register_default_workflows(self): """Register default workflow patterns.""" # Investigation workflow investigation_workflow = WorkflowDefinition( workflow_id="default_investigation", name="Standard Investigation Workflow", pattern=OrchestrationPattern.SEQUENTIAL, steps=[ WorkflowStep( step_id="anomaly_detection", agent_name="zumbi", action="detect_anomalies" ), WorkflowStep( step_id="pattern_analysis", agent_name="anita", action="analyze_patterns", conditions={"if": "anomalies_found", "gt": 0} ), WorkflowStep( step_id="report_generation", agent_name="tiradentes", action="generate_report" ) ] ) self._workflows["default_investigation"] = investigation_workflow async def execute_workflow( self, workflow_id: str, initial_data: Dict[str, Any], context: AgentContext ) -> Dict[str, Any]: """Execute a complete workflow.""" workflow = self._workflows.get(workflow_id) if not workflow: raise OrchestrationError(f"Workflow {workflow_id} not found") self.logger.info(f"Executing workflow: {workflow.name}") start_time = time.time() try: # Track metrics self._metrics.total_executions += 1 self._metrics.pattern_usage[workflow.pattern.value] += 1 # Execute based on pattern if workflow.pattern == OrchestrationPattern.SEQUENTIAL: result = await self._execute_sequential(workflow, initial_data, context) elif workflow.pattern == OrchestrationPattern.PARALLEL: result = await self._execute_parallel(workflow, initial_data, context) elif workflow.pattern == OrchestrationPattern.FAN_OUT_FAN_IN: result = await self._execute_fan_out_fan_in(workflow, initial_data, context) elif workflow.pattern == OrchestrationPattern.CONDITIONAL: result = await self._execute_conditional(workflow, initial_data, context) elif workflow.pattern == OrchestrationPattern.SAGA: result = await self._execute_saga(workflow, initial_data, context) elif workflow.pattern == OrchestrationPattern.MAP_REDUCE: result = await self._execute_map_reduce(workflow, initial_data, context) else: raise OrchestrationError(f"Unsupported pattern: {workflow.pattern}") # Update metrics duration = time.time() - start_time self._metrics.successful_executions += 1 self._metrics.total_duration_seconds += duration return { "workflow_id": workflow_id, "status": "completed", "result": result, "duration": duration } except Exception as e: self._metrics.failed_executions += 1 raise OrchestrationError(f"Workflow execution failed: {e}") async def _execute_sequential( self, workflow: WorkflowDefinition, data: Dict[str, Any], context: AgentContext ) -> Dict[str, Any]: """Execute workflow steps sequentially.""" current_data = data.copy() results = [] for step in workflow.steps: # Check conditions if not self._check_conditions(step.conditions, current_data): continue # Execute step step_result = await self._execute_step(step, current_data, context) results.append(step_result) # Map output to next input for output_key, data_key in step.output_mapping.items(): if output_key in step_result.get("data", {}): current_data[data_key] = step_result["data"][output_key] return { "pattern": "sequential", "steps_executed": len(results), "final_data": current_data, "step_results": results } async def _execute_parallel( self, workflow: WorkflowDefinition, data: Dict[str, Any], context: AgentContext ) -> Dict[str, Any]: """Execute workflow steps in parallel.""" tasks = [] for step in workflow.steps: if self._check_conditions(step.conditions, data): tasks.append(self._execute_step(step, data.copy(), context)) results = await asyncio.gather(*tasks, return_exceptions=True) return { "pattern": "parallel", "steps_executed": len(results), "results": [r for r in results if not isinstance(r, Exception)], "errors": [str(r) for r in results if isinstance(r, Exception)] } async def _execute_step( self, step: WorkflowStep, data: Dict[str, Any], context: AgentContext ) -> Dict[str, Any]: """Execute a single workflow step.""" start_time = time.time() try: # Get agent agent = await self._get_agent_with_circuit_breaker(step.agent_name) # Prepare input data input_data = {} for input_key, data_key in step.input_mapping.items(): if data_key in data: input_data[input_key] = data[data_key] # Create message message = AgentMessage( type=step.action, data=input_data or data, sender="orchestrator", metadata={"workflow_step": step.step_id} ) # Execute with timeout response = await asyncio.wait_for( agent.process(message, context), timeout=step.timeout ) # Track metrics duration = time.time() - start_time self._metrics.agent_execution_times[step.agent_name].append(duration) return { "step_id": step.step_id, "agent": step.agent_name, "success": response.success, "data": response.data, "duration": duration } except asyncio.TimeoutError: raise OrchestrationError(f"Step {step.step_id} timed out after {step.timeout}s") except Exception as e: raise OrchestrationError(f"Step {step.step_id} failed: {e}") def _check_conditions(self, conditions: Dict[str, Any], data: Dict[str, Any]) -> bool: """Check if conditions are met for step execution.""" if not conditions: return True # Simple condition evaluation if "if" in conditions: field = conditions["if"] if field not in data: return False value = data[field] if "eq" in conditions: return value == conditions["eq"] elif "gt" in conditions: return value > conditions["gt"] elif "lt" in conditions: return value < conditions["lt"] elif "in" in conditions: return value in conditions["in"] return True async def _get_agent_with_circuit_breaker(self, agent_name: str) -> BaseAgent: """Get agent with circuit breaker protection.""" circuit_breaker = self._circuit_breakers.get(agent_name) if not circuit_breaker: circuit_breaker = CircuitBreaker() self._circuit_breakers[agent_name] = circuit_breaker # Check circuit state if circuit_breaker.state == CircuitState.OPEN: # Check if recovery timeout has passed if (datetime.utcnow() - circuit_breaker.last_failure_time).seconds > circuit_breaker.recovery_timeout: circuit_breaker.state = CircuitState.HALF_OPEN circuit_breaker.success_count = 0 else: raise OrchestrationError(f"Circuit breaker open for {agent_name}") try: agent = await agent_lazy_loader.get_agent(agent_name) # Reset on success if circuit_breaker.state == CircuitState.HALF_OPEN: circuit_breaker.success_count += 1 if circuit_breaker.success_count >= circuit_breaker.half_open_requests: circuit_breaker.state = CircuitState.CLOSED circuit_breaker.failure_count = 0 return agent except Exception as e: # Update failure count circuit_breaker.failure_count += 1 circuit_breaker.last_failure_time = datetime.utcnow() if circuit_breaker.failure_count >= circuit_breaker.failure_threshold: circuit_breaker.state = CircuitState.OPEN raise async def select_best_agent(self, required_capabilities: List[str]) -> Optional[BaseAgent]: """Select the best agent based on required capabilities.""" best_match = None best_score = 0 for agent_name, capabilities in self._agent_capabilities.items(): # Calculate capability match score score = sum(1 for cap in required_capabilities if cap in capabilities) if score > best_score: best_score = score best_match = agent_name if best_match: return await agent_lazy_loader.get_agent(best_match) return None async def execute_with_retry( self, agent: BaseAgent, message: AgentMessage, context: AgentContext, max_retries: int = 3, backoff_multiplier: float = 2.0, fallback_agent: Optional[BaseAgent] = None ) -> AgentResponse: """Execute agent with retry logic and optional fallback.""" last_error = None for attempt in range(max_retries + 1): try: return await agent.process(message, context) except Exception as e: last_error = e if attempt < max_retries: wait_time = (backoff_multiplier ** attempt) * 1.0 await asyncio.sleep(wait_time) continue # Try fallback agent if available if fallback_agent: self.logger.warning(f"Primary agent failed, trying fallback") return await fallback_agent.process(message, context) raise raise OrchestrationError(f"All retry attempts failed: {last_error}") def configure_retry_policy(self, policy: Dict[str, Any]): """Configure global retry policy.""" self._retry_policy = policy def configure_circuit_breaker(self, config: Dict[str, Any]): """Configure circuit breaker settings.""" self._circuit_breaker_config = config async def execute_conditional_workflow( self, workflow_def: Dict[str, Any], initial_data: Dict[str, Any], context: AgentContext ) -> List[Dict[str, Any]]: """Execute a conditional workflow with branching.""" execution_path = [] current_step = workflow_def["start"] current_data = initial_data.copy() while current_step: step_def = workflow_def["steps"][current_step] # Execute step agent_name = step_def["agent"] agent = await agent_lazy_loader.get_agent(agent_name) message = AgentMessage( type="process", data=current_data, sender="conditional_workflow", metadata={"step": current_step} ) response = await agent.process(message, context) execution_path.append({ "step": current_step, "agent": agent_name, "success": response.success, "data": response.data }) # Determine next step based on conditions next_step_def = step_def.get("next") if not next_step_def: break if isinstance(next_step_def, str): current_step = next_step_def else: # Conditional branching condition = next_step_def.get("condition") if condition == "anomalies_found": if response.data.get("anomalies_detected", 0) > 0: current_step = next_step_def.get("true") else: current_step = next_step_def.get("false") elif condition == "high_risk": if response.data.get("risk_level") == "high": current_step = next_step_def.get("true") else: current_step = next_step_def.get("false") else: current_step = next_step_def.get("default") # Update data for next step current_data.update(response.data) return execution_path async def discover_agents(self) -> List[Dict[str, Any]]: """Discover all available agents.""" return await agent_lazy_loader.list_agents() async def find_agents_with_capability(self, capability: str) -> List[Dict[str, Any]]: """Find agents with a specific capability.""" matching_agents = [] for agent_name, capabilities in self._agent_capabilities.items(): if capability in capabilities: agent_info = { "name": agent_name, "capabilities": capabilities } matching_agents.append(agent_info) return matching_agents async def execute_with_circuit_breaker( self, agent: BaseAgent, message: AgentMessage, context: AgentContext ) -> AgentResponse: """Execute agent with circuit breaker protection.""" agent_name = agent.name circuit_breaker = self._circuit_breakers.get(agent_name) if not circuit_breaker: circuit_breaker = CircuitBreaker(**self._circuit_breaker_config) self._circuit_breakers[agent_name] = circuit_breaker if circuit_breaker.state == CircuitState.OPEN: if (datetime.utcnow() - circuit_breaker.last_failure_time).seconds < circuit_breaker.recovery_timeout: raise OrchestrationError(f"Circuit breaker open for {agent_name}") else: circuit_breaker.state = CircuitState.HALF_OPEN try: response = await agent.process(message, context) if circuit_breaker.state == CircuitState.HALF_OPEN: circuit_breaker.success_count += 1 if circuit_breaker.success_count >= circuit_breaker.half_open_requests: circuit_breaker.state = CircuitState.CLOSED circuit_breaker.failure_count = 0 return response except Exception as e: circuit_breaker.failure_count += 1 circuit_breaker.last_failure_time = datetime.utcnow() if circuit_breaker.failure_count >= circuit_breaker.failure_threshold: circuit_breaker.state = CircuitState.OPEN raise async def start_saga( self, saga_definition: Dict[str, Any], initial_data: Dict[str, Any], context: AgentContext ) -> Dict[str, Any]: """Start a saga transaction.""" saga_state = { "saga_id": f"saga_{datetime.utcnow().timestamp()}", "name": saga_definition["name"], "current_step": 0, "completed_steps": [], "compensated_steps": [], "data": initial_data, "completed": False, "failed": False } return saga_state async def execute_saga_step( self, saga_state: Dict[str, Any], step: Dict[str, Any], context: AgentContext ) -> Dict[str, Any]: """Execute a saga step with compensation support.""" try: agent_name = step["agent"] agent = await agent_lazy_loader.get_agent(agent_name) message = AgentMessage( type=step["service"], data=saga_state["data"], sender="saga_coordinator", metadata={"saga_id": saga_state["saga_id"]} ) response = await agent.process(message, context) if response.success: saga_state["completed_steps"].append({ "step": step, "result": response.data }) saga_state["current_step"] += 1 if saga_state["current_step"] >= len(saga_state.get("total_steps", [])): saga_state["completed"] = True else: saga_state["failed"] = True # Trigger compensation await self._compensate_saga(saga_state, context) except Exception as e: saga_state["failed"] = True saga_state["error"] = str(e) await self._compensate_saga(saga_state, context) return saga_state async def _compensate_saga(self, saga_state: Dict[str, Any], context: AgentContext): """Compensate completed saga steps.""" for completed_step in reversed(saga_state["completed_steps"]): try: step = completed_step["step"] if "compensation" in step: agent = await agent_lazy_loader.get_agent(step["agent"]) compensation_message = AgentMessage( type=step["compensation"], data={ "original_data": saga_state["data"], "step_result": completed_step["result"] }, sender="saga_compensator", metadata={"saga_id": saga_state["saga_id"]} ) await agent.process(compensation_message, context) saga_state["compensated_steps"].append(step) except Exception as e: self.logger.error(f"Compensation failed for step: {e}") def get_event_bus(self) -> EventBus: """Get the event bus for choreography.""" return self._event_bus async def get_stats(self) -> Dict[str, Any]: """Get orchestrator statistics.""" return { "total_executions": self._metrics.total_executions, "successful_executions": self._metrics.successful_executions, "failed_executions": self._metrics.failed_executions, "success_rate": ( self._metrics.successful_executions / self._metrics.total_executions if self._metrics.total_executions > 0 else 0 ), "average_duration": ( self._metrics.total_duration_seconds / self._metrics.successful_executions if self._metrics.successful_executions > 0 else 0 ), "pattern_usage": dict(self._metrics.pattern_usage), "agent_performance": { agent: { "executions": len(times), "avg_time": sum(times) / len(times) if times else 0, "min_time": min(times) if times else 0, "max_time": max(times) if times else 0 } for agent, times in self._metrics.agent_execution_times.items() }, "circuit_breakers": { agent: { "state": cb.state.value, "failure_count": cb.failure_count } for agent, cb in self._circuit_breakers.items() } } async def _execute_fan_out_fan_in( self, workflow: WorkflowDefinition, data: Dict[str, Any], context: AgentContext ) -> Dict[str, Any]: """Execute fan-out/fan-in pattern.""" # Fan-out: execute multiple steps in parallel fan_out_results = await self._execute_parallel(workflow, data, context) # Fan-in: aggregate results aggregated_data = { "pattern": "fan_out_fan_in", "fan_out_results": fan_out_results["results"], "aggregated_data": {} } # Simple aggregation - can be customized for result in fan_out_results["results"]: if result.get("success") and "data" in result: aggregated_data["aggregated_data"].update(result["data"]) return aggregated_data async def _execute_map_reduce( self, workflow: WorkflowDefinition, data: Dict[str, Any], context: AgentContext ) -> Dict[str, Any]: """Execute map-reduce pattern.""" # Map phase map_items = data.get("items", []) map_results = [] for item in map_items: # Execute map step for each item map_step = workflow.steps[0] # Assuming first step is map step_result = await self._execute_step( map_step, {"item": item}, context ) map_results.append(step_result) # Reduce phase reduce_step = workflow.steps[1] # Assuming second step is reduce reduce_data = { "map_results": [r["data"] for r in map_results if r.get("success")] } reduce_result = await self._execute_step( reduce_step, reduce_data, context ) return { "pattern": "map_reduce", "map_count": len(map_results), "reduce_result": reduce_result["data"] if reduce_result.get("success") else None } async def _execute_saga( self, workflow: WorkflowDefinition, data: Dict[str, Any], context: AgentContext ) -> Dict[str, Any]: """Execute saga pattern with compensation.""" saga_state = await self.start_saga( {"name": workflow.name, "steps": workflow.steps}, data, context ) saga_state["total_steps"] = workflow.steps for step in workflow.steps: saga_state = await self.execute_saga_step( saga_state, {"agent": step.agent_name, "service": step.action}, context ) if saga_state["failed"]: break return { "pattern": "saga", "saga_id": saga_state["saga_id"], "completed": saga_state["completed"], "failed": saga_state["failed"], "completed_steps": len(saga_state["completed_steps"]), "compensated_steps": len(saga_state["compensated_steps"]) } # Global orchestrator instance orchestrator = AgentOrchestrator() async def get_orchestrator() -> AgentOrchestrator: """Get the global orchestrator instance.""" return orchestrator