src/agents/supervisor_agent.py

"""
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)
        }