""" Supervisor agent for orchestrating task execution using LangGraph. """ import time from typing import Dict, Any, List, Optional, Callable from dataclasses import dataclass from enum import Enum try: from langgraph.graph import StateGraph, END from langgraph.checkpoint.memory import MemorySaver LANGGRAPH_AVAILABLE = True except ImportError: LANGGRAPH_AVAILABLE = False from .base_agent import BaseAgent, AgentResult, AgentState, AgentMessage from src.utils.logging_config import logger class TaskStatus(Enum): """Task execution status.""" PENDING = "pending" IN_PROGRESS = "in_progress" COMPLETED = "completed" FAILED = "failed" RETRYING = "retrying" @dataclass class Task: """Represents a task to be executed.""" id: str name: str agent_type: str input_data: Dict[str, Any] dependencies: List[str] = None status: TaskStatus = TaskStatus.PENDING result: Optional[AgentResult] = None retry_count: int = 0 max_retries: int = 3 @dataclass class WorkflowState: """State of the workflow execution.""" query: str tasks: List[Task] results: Dict[str, Any] current_task: Optional[str] = None error_message: Optional[str] = None completed: bool = False metadata: Dict[str, Any] = None class SupervisorAgent(BaseAgent): """Supervisor agent that orchestrates task execution using LangGraph.""" def __init__(self, agent_id: Optional[str] = None, config: Optional[Dict[str, Any]] = None): """Initialize the supervisor agent.""" if not LANGGRAPH_AVAILABLE: raise ImportError("LangGraph is required for SupervisorAgent. Install with: pip install langgraph") super().__init__(agent_id, config) # Worker agents registry self.worker_agents: Dict[str, BaseAgent] = {} # Workflow configuration self.max_retries = self.config.get('max_retries', 3) self.retry_delay = self.config.get('retry_delay', 1.0) # Build the workflow graph self._build_workflow_graph() def _initialize(self) -> None: """Initialize supervisor-specific components.""" self.checkpointer = MemorySaver() logger.info("Supervisor agent initialized with LangGraph") def _build_workflow_graph(self) -> None: """Build the LangGraph workflow.""" # Create the state graph workflow = StateGraph(WorkflowState) # Add nodes for each step workflow.add_node("parse_query", self._parse_query_node) workflow.add_node("match_apis", self._match_apis_node) workflow.add_node("execute_apis", self._execute_apis_node) workflow.add_node("format_results", self._format_results_node) workflow.add_node("evaluate_results", self._evaluate_results_node) # Define the workflow edges workflow.set_entry_point("parse_query") workflow.add_edge("parse_query", "match_apis") workflow.add_edge("match_apis", "execute_apis") workflow.add_edge("execute_apis", "format_results") workflow.add_edge("format_results", "evaluate_results") workflow.add_edge("evaluate_results", END) # Compile the graph self.workflow = workflow.compile(checkpointer=self.checkpointer) logger.info("LangGraph workflow compiled successfully") def register_worker(self, agent_type: str, agent: BaseAgent) -> None: """ Register a worker agent. Args: agent_type: Type identifier for the agent agent: Worker agent instance """ self.worker_agents[agent_type] = agent logger.info(f"Registered worker agent: {agent_type} ({agent.agent_id})") def execute(self, input_data: Dict[str, Any]) -> AgentResult: """ Execute the supervised workflow. Args: input_data: Input containing query and configuration Returns: AgentResult with workflow outcome """ start_time = time.time() self.set_state(AgentState.RUNNING) try: # Validate input if not self.validate_input(input_data): raise ValueError("Invalid input data") query = input_data.get('query', '') if not query: raise ValueError("Query is required") # Initialize workflow state initial_state = WorkflowState( query=query, tasks=[], results={}, metadata=input_data.get('metadata', {}) ) # Execute the workflow config = {"configurable": {"thread_id": f"workflow_{int(time.time())}"}} final_state = self.workflow.invoke(initial_state, config) execution_time = time.time() - start_time if final_state.completed and not final_state.error_message: self.set_state(AgentState.COMPLETED) result = AgentResult( agent_id=self.agent_id, success=True, data=final_state.results, execution_time=execution_time, metadata={ 'tasks_completed': len([t for t in final_state.tasks if t.status == TaskStatus.COMPLETED]), 'workflow_state': final_state } ) else: self.set_state(AgentState.FAILED) result = AgentResult( agent_id=self.agent_id, success=False, data=final_state.results, error_message=final_state.error_message or "Workflow failed", execution_time=execution_time, metadata={'workflow_state': final_state} ) self.log_execution(result) return result except Exception as e: execution_time = time.time() - start_time self.set_state(AgentState.FAILED) result = AgentResult( agent_id=self.agent_id, success=False, data=None, error_message=str(e), execution_time=execution_time ) self.log_execution(result) return result def _parse_query_node(self, state: WorkflowState) -> WorkflowState: """Parse the input query.""" logger.info(f"Parsing query: {state.query}") try: # Execute query parsing using worker agent parser_agent = self.worker_agents.get('query_parser') if not parser_agent: raise ValueError("Query parser agent not registered") parse_result = parser_agent.execute({'query': state.query}) if parse_result.success: state.results['parsed_query'] = parse_result.data logger.info("Query parsing completed successfully") else: state.error_message = f"Query parsing failed: {parse_result.error_message}" logger.error(state.error_message) except Exception as e: state.error_message = f"Query parsing error: {str(e)}" logger.error(state.error_message) return state def _match_apis_node(self, state: WorkflowState) -> WorkflowState: """Match parsed query to available APIs.""" logger.info("Matching APIs") if state.error_message: return state try: # Get parsed query results parsed_query = state.results.get('parsed_query') if not parsed_query: raise ValueError("No parsed query available") # Execute API matching matcher_agent = self.worker_agents.get('api_matcher') if matcher_agent: match_result = matcher_agent.execute({ 'keywords': parsed_query.get('keywords', []), 'intent': parsed_query.get('intent') }) if match_result.success: state.results['api_matches'] = match_result.data logger.info(f"Found {len(match_result.data.get('matches', []))} API matches") else: logger.warning(f"API matching failed: {match_result.error_message}") state.results['api_matches'] = {'matches': []} else: logger.warning("API matcher agent not registered") state.results['api_matches'] = {'matches': []} except Exception as e: state.error_message = f"API matching error: {str(e)}" logger.error(state.error_message) return state def _execute_apis_node(self, state: WorkflowState) -> WorkflowState: """Execute matched API calls.""" logger.info("Executing API calls") if state.error_message: return state try: # Get API matches api_matches = state.results.get('api_matches', {}) matches = api_matches.get('matches', []) if not matches: logger.info("No API matches to execute") state.results['api_results'] = [] return state # Execute API calls using executor agent executor_agent = self.worker_agents.get('api_executor') if not executor_agent: raise ValueError("API executor agent not registered") execution_result = executor_agent.execute({ 'matches': matches, 'query_context': state.results.get('parsed_query') }) if execution_result.success: state.results['api_results'] = execution_result.data logger.info(f"Executed {len(execution_result.data.get('results', []))} API calls") else: state.error_message = f"API execution failed: {execution_result.error_message}" logger.error(state.error_message) except Exception as e: state.error_message = f"API execution error: {str(e)}" logger.error(state.error_message) return state def _format_results_node(self, state: WorkflowState) -> WorkflowState: """Format the results for output.""" logger.info("Formatting results") if state.error_message: return state try: # Format results using formatter agent formatter_agent = self.worker_agents.get('result_formatter') if not formatter_agent: # Basic formatting if no formatter agent state.results['formatted_output'] = { 'query': state.query, 'results': state.results.get('api_results', []), 'metadata': state.metadata } logger.info("Applied basic result formatting") return state format_result = formatter_agent.execute({ 'query': state.query, 'parsed_query': state.results.get('parsed_query'), 'api_matches': state.results.get('api_matches'), 'api_results': state.results.get('api_results'), 'metadata': state.metadata }) if format_result.success: state.results['formatted_output'] = format_result.data logger.info("Result formatting completed successfully") else: state.error_message = f"Result formatting failed: {format_result.error_message}" logger.error(state.error_message) except Exception as e: state.error_message = f"Result formatting error: {str(e)}" logger.error(state.error_message) return state def _evaluate_results_node(self, state: WorkflowState) -> WorkflowState: """Evaluate the final results.""" logger.info("Evaluating results") try: # Basic evaluation - can be enhanced with evaluator agent evaluator_agent = self.worker_agents.get('evaluator') if evaluator_agent: eval_result = evaluator_agent.execute({ 'query': state.query, 'results': state.results, 'workflow_state': state }) if eval_result.success: state.results['evaluation'] = eval_result.data # Check if results meet quality threshold quality_score = eval_result.data.get('quality_score', 0.5) if quality_score >= 0.7: state.completed = True logger.info(f"Workflow completed successfully (quality: {quality_score:.2f})") else: logger.warning(f"Results quality below threshold: {quality_score:.2f}") state.completed = True # Complete anyway for now else: logger.warning(f"Result evaluation failed: {eval_result.error_message}") state.completed = True # Complete anyway else: # Basic evaluation without evaluator agent has_results = bool(state.results.get('api_results')) state.completed = True state.results['evaluation'] = { 'has_results': has_results, 'quality_score': 0.8 if has_results else 0.3, 'evaluation_method': 'basic' } logger.info(f"Basic evaluation completed: {'success' if has_results else 'limited results'}") except Exception as e: state.error_message = f"Result evaluation error: {str(e)}" logger.error(state.error_message) state.completed = True # Complete with error return state def get_capabilities(self) -> List[str]: """Get supervisor capabilities.""" return [ "workflow_orchestration", "task_scheduling", "agent_coordination", "error_handling", "result_aggregation" ] def get_workflow_status(self) -> Dict[str, Any]: """Get current workflow status.""" return { 'registered_workers': list(self.worker_agents.keys()), 'workflow_available': LANGGRAPH_AVAILABLE, 'execution_history_count': len(self.execution_history) }