src/agents/base_agent.py

"""
Base agent class for the agentic AI system.
"""

from abc import ABC, abstractmethod
from typing import Dict, Any, Optional, List
from dataclasses import dataclass
from enum import Enum
import uuid
from datetime import datetime

from src.utils.logging_config import logger


class AgentState(Enum):
    """Agent execution states."""
    IDLE = "idle"
    RUNNING = "running"
    COMPLETED = "completed"
    FAILED = "failed"
    PAUSED = "paused"


@dataclass
class AgentMessage:
    """Message passed between agents."""
    id: str
    sender: str
    recipient: str
    message_type: str
    content: Dict[str, Any]
    timestamp: datetime
    metadata: Dict[str, Any] = None


@dataclass
class AgentResult:
    """Result from agent execution."""
    agent_id: str
    success: bool
    data: Any
    error_message: Optional[str] = None
    execution_time: float = 0.0
    metadata: Dict[str, Any] = None


class BaseAgent(ABC):
    """Base class for all agents in the system."""
    
    def __init__(self, agent_id: Optional[str] = None, config: Optional[Dict[str, Any]] = None):
        """
        Initialize the agent.
        
        Args:
            agent_id: Unique identifier for the agent
            config: Agent configuration
        """
        self.agent_id = agent_id or f"{self.__class__.__name__}_{uuid.uuid4().hex[:8]}"
        self.config = config or {}
        self.state = AgentState.IDLE
        self.message_history: List[AgentMessage] = []
        self.execution_history: List[AgentResult] = []
        
        # Initialize agent-specific components
        self._initialize()
        
        logger.info(f"Initialized agent: {self.agent_id}")
    
    @abstractmethod
    def _initialize(self) -> None:
        """Initialize agent-specific components."""
        pass
    
    @abstractmethod
    def execute(self, input_data: Dict[str, Any]) -> AgentResult:
        """
        Execute the agent's main task.
        
        Args:
            input_data: Input data for execution
            
        Returns:
            AgentResult with execution outcome
        """
        pass
    
    def send_message(self, recipient: str, message_type: str, content: Dict[str, Any]) -> AgentMessage:
        """
        Send a message to another agent.
        
        Args:
            recipient: Target agent ID
            message_type: Type of message
            content: Message content
            
        Returns:
            AgentMessage that was sent
        """
        message = AgentMessage(
            id=uuid.uuid4().hex,
            sender=self.agent_id,
            recipient=recipient,
            message_type=message_type,
            content=content,
            timestamp=datetime.utcnow()
        )
        
        self.message_history.append(message)
        logger.debug(f"Agent {self.agent_id} sent message to {recipient}: {message_type}")
        
        return message
    
    def receive_message(self, message: AgentMessage) -> Optional[AgentMessage]:
        """
        Receive and process a message.
        
        Args:
            message: Incoming message
            
        Returns:
            Optional response message
        """
        self.message_history.append(message)
        logger.debug(f"Agent {self.agent_id} received message from {message.sender}: {message.message_type}")
        
        # Process message based on type
        return self._process_message(message)
    
    def _process_message(self, message: AgentMessage) -> Optional[AgentMessage]:
        """
        Process an incoming message.
        
        Args:
            message: Message to process
            
        Returns:
            Optional response message
        """
        # Default implementation - subclasses should override
        logger.debug(f"Agent {self.agent_id} processing message: {message.message_type}")
        return None
    
    def get_state(self) -> AgentState:
        """Get current agent state."""
        return self.state
    
    def set_state(self, state: AgentState) -> None:
        """Set agent state."""
        old_state = self.state
        self.state = state
        logger.debug(f"Agent {self.agent_id} state changed: {old_state} -> {state}")
    
    def get_capabilities(self) -> List[str]:
        """
        Get list of agent capabilities.
        
        Returns:
            List of capability names
        """
        return []
    
    def get_status(self) -> Dict[str, Any]:
        """
        Get agent status information.
        
        Returns:
            Status dictionary
        """
        return {
            'agent_id': self.agent_id,
            'agent_type': self.__class__.__name__,
            'state': self.state.value,
            'capabilities': self.get_capabilities(),
            'message_count': len(self.message_history),
            'execution_count': len(self.execution_history),
            'last_execution': self.execution_history[-1].timestamp if self.execution_history else None
        }
    
    def reset(self) -> None:
        """Reset agent to initial state."""
        self.state = AgentState.IDLE
        self.message_history.clear()
        self.execution_history.clear()
        logger.info(f"Agent {self.agent_id} reset")
    
    def validate_input(self, input_data: Dict[str, Any]) -> bool:
        """
        Validate input data.
        
        Args:
            input_data: Data to validate
            
        Returns:
            True if valid, False otherwise
        """
        # Basic validation - subclasses should override
        return isinstance(input_data, dict)
    
    def log_execution(self, result: AgentResult) -> None:
        """Log execution result."""
        self.execution_history.append(result)
        
        if result.success:
            logger.info(f"Agent {self.agent_id} execution completed successfully in {result.execution_time:.2f}s")
        else:
            logger.error(f"Agent {self.agent_id} execution failed: {result.error_message}")
    
    def __str__(self) -> str:
        """String representation of the agent."""
        return f"{self.__class__.__name__}(id={self.agent_id}, state={self.state.value})"
    
    def __repr__(self) -> str:
        """Detailed string representation."""
        return f"{self.__class__.__name__}(id='{self.agent_id}', state={self.state}, config={self.config})"