""" 🤖 Agentic AutoML - Base Agent Framework This module defines the foundational components for the multi-agent system: - Message types for agent communication - Base Agent class with lifecycle methods - Agent result and status types """ from abc import ABC, abstractmethod from dataclasses import dataclass, field from enum import Enum from typing import Any, Dict, List, Optional, Callable from datetime import datetime from agents.memory import AgentMemory import uuid import logging logger = logging.getLogger(__name__) # ============================================================================= # MESSAGE TYPES - For inter-agent communication # ============================================================================= class MessageType(Enum): """Types of messages agents can send""" # Decision messages APPROVAL = "approval" REJECTION = "rejection" RETRY = "retry" # Signal messages ESCALATE = "escalate" REQUEST_FEEDBACK = "request_feedback" PHASE_COMPLETE = "phase_complete" # Control messages ABORT = "abort" PAUSE = "pause" RESUME = "resume" class AgentStatus(Enum): """Agent execution status""" IDLE = "idle" RUNNING = "running" SUCCESS = "success" FAILED = "failed" WAITING = "waiting" RETRY = "retry" class Phase(Enum): """Pipeline phases""" FAST_DISCOVERY = "fast_discovery" DEEP_VALIDATION = "deep_validation" @dataclass class AgentMessage: """Message passed between agents""" id: str = field(default_factory=lambda: str(uuid.uuid4())[:8]) sender: str = "" receiver: str = "" type: MessageType = MessageType.APPROVAL payload: Dict[str, Any] = field(default_factory=dict) timestamp: datetime = field(default_factory=datetime.now) def to_dict(self) -> Dict: return { "id": self.id, "sender": self.sender, "receiver": self.receiver, "type": self.type.value, "payload": self.payload, "timestamp": self.timestamp.isoformat() } @dataclass class AgentResult: """Result returned by an agent after execution""" status: AgentStatus agent_name: str phase: Phase data: Dict[str, Any] = field(default_factory=dict) messages: List[AgentMessage] = field(default_factory=list) errors: List[str] = field(default_factory=list) recommendations: List[str] = field(default_factory=list) metrics: Dict[str, float] = field(default_factory=dict) duration_seconds: float = 0.0 @property def success(self) -> bool: return self.status == AgentStatus.SUCCESS @property def should_retry(self) -> bool: return self.status == AgentStatus.RETRY def add_message(self, receiver: str, msg_type: MessageType, payload: Dict = None): """Helper to add a message to send""" self.messages.append(AgentMessage( sender=self.agent_name, receiver=receiver, type=msg_type, payload=payload or {} )) # ============================================================================= # BASE AGENT CLASS # ============================================================================= class BaseAgent(ABC): """ Base class for all agents in the Agentic AutoML system. Each agent has: - A name and description - Access to shared memory - Ability to send/receive messages - Lifecycle methods (validate, execute, handle_failure) """ name: str = "BaseAgent" description: str = "Base agent class" def __init__(self, memory: 'AgentMemory' = None): self.memory = memory self.current_phase = Phase.FAST_DISCOVERY self.retry_count = 0 self.max_retries = 3 self._start_time = None self.logger = logging.getLogger(f"agent.{self.name}") # ========================================================================= # LIFECYCLE METHODS # ========================================================================= def run(self, **kwargs) -> AgentResult: """ Main entry point for agent execution. Handles timing, error catching, and retry logic. """ import time self._start_time = time.time() try: # Pre-execution validation validation = self.validate(**kwargs) if not validation['valid']: return AgentResult( status=AgentStatus.FAILED, agent_name=self.name, phase=self.current_phase, errors=validation.get('errors', ['Validation failed']), duration_seconds=time.time() - self._start_time ) # Main execution self.logger.info(f"🤖 {self.name} starting...") result = self.execute(**kwargs) result.duration_seconds = time.time() - self._start_time # Log result if result.success: self.logger.info(f"✅ {self.name} completed in {result.duration_seconds:.2f}s") else: self.logger.warning(f"⚠️ {self.name} finished with status: {result.status.value}") return result except Exception as e: self.logger.error(f"❌ {self.name} error: {str(e)}") return self.handle_failure(e, **kwargs) def validate(self, **kwargs) -> Dict[str, Any]: """ Validate inputs before execution. Override in subclasses for specific validation. """ return {'valid': True} @abstractmethod def execute(self, **kwargs) -> AgentResult: """ Main execution logic. Must be implemented by subclasses. """ pass def handle_failure(self, error: Exception, **kwargs) -> AgentResult: """ Handle execution failures. Can be overridden for custom recovery. """ import time self.retry_count += 1 if self.retry_count < self.max_retries: return AgentResult( status=AgentStatus.RETRY, agent_name=self.name, phase=self.current_phase, errors=[str(error)], recommendations=[f"Retry attempt {self.retry_count}/{self.max_retries}"], duration_seconds=time.time() - self._start_time if self._start_time else 0 ) return AgentResult( status=AgentStatus.FAILED, agent_name=self.name, phase=self.current_phase, errors=[str(error), f"Max retries ({self.max_retries}) exceeded"], duration_seconds=time.time() - self._start_time if self._start_time else 0 ) # ========================================================================= # PHASE CONTROL # ========================================================================= def set_phase(self, phase: Phase): """Set the current execution phase""" self.current_phase = phase self.logger.info(f"📍 Phase set to: {phase.value}") def is_fast_phase(self) -> bool: """Check if in fast discovery phase""" return self.current_phase == Phase.FAST_DISCOVERY def is_deep_phase(self) -> bool: """Check if in deep validation phase""" return self.current_phase == Phase.DEEP_VALIDATION # ========================================================================= # MEMORY ACCESS # ========================================================================= def read_state(self, key: str, default: Any = None) -> Any: """Read from shared memory""" if self.memory: return self.memory.get(key, default) return default def write_state(self, key: str, value: Any, stage: str = None): """Write to shared memory (immutable per stage)""" if self.memory: self.memory.set(key, value, stage or self.name) def get_artifacts(self, artifact_type: str) -> List[Any]: """Get artifacts of a specific type""" if self.memory: return self.memory.get_artifacts(artifact_type) return [] # ============================================================================= # AGENT REGISTRY # ============================================================================= class AgentRegistry: """Registry for managing agent instances""" _agents: Dict[str, BaseAgent] = {} @classmethod def register(cls, agent: BaseAgent): """Register an agent""" cls._agents[agent.name] = agent @classmethod def get(cls, name: str) -> Optional[BaseAgent]: """Get an agent by name""" return cls._agents.get(name) @classmethod def all(cls) -> List[BaseAgent]: """Get all registered agents""" return list(cls._agents.values()) @classmethod def clear(cls): """Clear all registered agents""" cls._agents.clear()