""" 🤖 Agentic AutoML - Agent Orchestrator The central coordinator that: - Manages agent execution order - Handles two-phase optimization (Fast Discovery → Deep Validation) - Implements feedback loops and retry logic - Routes messages between agents """ from typing import Dict, List, Optional, Any, Callable from datetime import datetime from enum import Enum import logging import uuid import time from .base import ( BaseAgent, AgentResult, AgentStatus, AgentMessage, MessageType, Phase, AgentRegistry ) from .memory import AgentMemory logger = logging.getLogger(__name__) class PipelineStatus(Enum): """Overall pipeline status""" IDLE = "idle" RUNNING = "running" FAST_DISCOVERY = "fast_discovery" DEEP_VALIDATION = "deep_validation" SUCCESS = "success" FAILED = "failed" CANCELLED = "cancelled" class AgentOrchestrator: """ Central orchestrator for the Agentic AutoML pipeline. Responsibilities: 1. Coordinate agent execution order 2. Manage two-phase optimization strategy 3. Handle message routing between agents 4. Implement feedback loops for failure recovery 5. Track pipeline state and metrics """ def __init__(self): self.memory = AgentMemory() self.agents: Dict[str, BaseAgent] = {} self.status = PipelineStatus.IDLE self.current_phase = Phase.FAST_DISCOVERY self.iteration = 0 self.max_iterations = 5 self.message_queue: List[AgentMessage] = [] self.execution_log: List[Dict] = [] # Thresholds for phase transition self.fast_phase_threshold = 0.5 # Minimum score to enter deep phase self.approval_threshold = 0.7 # Minimum score for final approval # Callbacks self.on_progress: Optional[Callable[[str, float], None]] = None self.check_cancellation: Optional[Callable[[], None]] = None # Timing self.start_time: Optional[datetime] = None # ========================================================================= # AGENT REGISTRATION # ========================================================================= def register_agent(self, name: str, agent: BaseAgent): """Register an agent with the orchestrator""" agent.memory = self.memory self.agents[name] = agent logger.info(f"📝 Registered agent: {name}") def get_agent(self, name: str) -> Optional[BaseAgent]: """Get an agent by name""" return self.agents.get(name) # ========================================================================= # PIPELINE EXECUTION # ========================================================================= def run(self, dataset, target_column: str, task_type: str = None, on_progress: Callable[[str, float], None] = None, check_cancellation: Callable[[], None] = None) -> Dict[str, Any]: """ Run the complete agentic AutoML pipeline. Args: dataset: Input DataFrame target_column: Name of target column task_type: 'classification' or 'regression' (auto-detected if None) on_progress: Callback for progress updates check_cancellation: Callback to check for user cancellation Returns: Dict with results, metrics, and artifacts """ self.start_time = datetime.now() self.on_progress = on_progress self.check_cancellation = check_cancellation self.status = PipelineStatus.RUNNING self.iteration = 0 # Initialize memory self.memory.clear() self.memory.pipeline_id = str(uuid.uuid4())[:8] self.memory.set("dataset", dataset, "orchestrator") self.memory.set("target_column", target_column, "orchestrator") self.memory.set("task_type", task_type, "orchestrator") logger.info(f"🚀 Starting Agentic AutoML Pipeline (ID: {self.memory.pipeline_id})") self._report_progress("Initializing pipeline", 0.0) try: # ========== PHASE 1: FAST DISCOVERY ========== self.current_phase = Phase.FAST_DISCOVERY self.status = PipelineStatus.FAST_DISCOVERY logger.info("⚡ PHASE 1: FAST DISCOVERY") fast_result = self._run_fast_discovery() if not fast_result['success']: return self._build_failure_result(fast_result.get('errors', [])) # Check if we should proceed to deep validation fast_score = fast_result.get('score', 0) if fast_score < self.fast_phase_threshold: logger.warning(f"⚠️ Fast phase score ({fast_score:.2f}) below threshold") return self._build_result(fast_result, approved=False) # ========== PHASE 2: DEEP VALIDATION ========== self.current_phase = Phase.DEEP_VALIDATION self.status = PipelineStatus.DEEP_VALIDATION logger.info("🔍 PHASE 2: DEEP VALIDATION") deep_result = self._run_deep_validation() if not deep_result['success']: # Failed deep validation - return best fast phase result logger.warning("⚠️ Deep validation failed, using fast phase result") return self._build_result(fast_result, approved=False) # ========== FINAL APPROVAL ========== final_score = deep_result.get('score', 0) approved = final_score >= self.approval_threshold self.status = PipelineStatus.SUCCESS if approved else PipelineStatus.FAILED return self._build_result(deep_result, approved=approved) except Exception as e: logger.error(f"❌ Pipeline error: {str(e)}") self.status = PipelineStatus.FAILED return self._build_failure_result([str(e)]) # ========================================================================= # PHASE EXECUTION # ========================================================================= def _run_fast_discovery(self) -> Dict[str, Any]: """ Fast Discovery Phase: - Limited algorithms (top 3-4) - Shallow hyperparameter search - Quick feature screening - Single train-test split """ self._check_cancelled() # Set all agents to fast phase for agent in self.agents.values(): agent.set_phase(Phase.FAST_DISCOVERY) # Execute data pipeline self._report_progress("Data Quality Check", 0.1) data_result = self._run_agent("data_quality") if not data_result.success: return {"success": False, "errors": data_result.errors} self._report_progress("Preprocessing", 0.2) prep_result = self._run_agent("preprocessing") if not prep_result.success: return {"success": False, "errors": prep_result.errors} self._report_progress("Feature Engineering", 0.3) feature_result = self._run_agent("feature_engineer") if not feature_result.success: return {"success": False, "errors": feature_result.errors} # Execute model pipeline self._report_progress("Model Selection", 0.4) model_result = self._run_agent("model_strategy") if not model_result.success: return {"success": False, "errors": model_result.errors} self._report_progress("Hyperparameter Tuning", 0.5) hyperparam_result = self._run_agent("hyperparam") if not hyperparam_result.success: return {"success": False, "errors": hyperparam_result.errors} self._report_progress("Training Validation", 0.6) training_result = self._run_agent("training_validator") # Handle training validation result with feedback loop if training_result.should_retry: return self._handle_feedback_loop(training_result) if not training_result.success: return {"success": False, "errors": training_result.errors} return { "success": True, "score": training_result.metrics.get("score", 0), "metrics": training_result.metrics, "phase": "fast_discovery" } def _run_deep_validation(self) -> Dict[str, Any]: """ Deep Validation Phase: - Robust cross-validation - Stability checks - Overfitting detection - Feature ablation - Drift sensitivity """ self._check_cancelled() # Set all agents to deep phase for agent in self.agents.values(): agent.set_phase(Phase.DEEP_VALIDATION) self._report_progress("Evaluation & Generalization", 0.7) eval_result = self._run_agent("evaluation") if eval_result.should_retry: return self._handle_feedback_loop(eval_result) if not eval_result.success: return {"success": False, "errors": eval_result.errors} # Only generate visualizations if evaluation passed self._report_progress("Generating Explanations", 0.85) viz_result = self._run_agent("visualization") # Prepare for deployment self._report_progress("Deployment Preparation", 0.95) deploy_result = self._run_agent("deployment") return { "success": True, "score": eval_result.metrics.get("score", 0), "metrics": eval_result.metrics, "phase": "deep_validation", "approved": deploy_result.success } # ========================================================================= # FEEDBACK LOOPS # ========================================================================= def _handle_feedback_loop(self, result: AgentResult) -> Dict[str, Any]: """ Handle feedback from validators. Routes back to appropriate agent based on recommendations. """ self.iteration += 1 if self.iteration >= self.max_iterations: logger.warning(f"⚠️ Max iterations ({self.max_iterations}) reached") return {"success": False, "errors": ["Max iterations reached"]} logger.info(f"🔄 Feedback loop iteration {self.iteration}") # Parse recommendations to determine which agent to retry for msg in result.messages: if msg.type == MessageType.RETRY: target_agent = msg.receiver logger.info(f" → Routing to {target_agent}") # Re-run from that agent retry_result = self._run_agent(target_agent) if retry_result.success: # Continue from where we left off if self.current_phase == Phase.FAST_DISCOVERY: return self._run_fast_discovery() else: return self._run_deep_validation() # Default: retry current phase if self.current_phase == Phase.FAST_DISCOVERY: return self._run_fast_discovery() else: return self._run_deep_validation() # ========================================================================= # HELPER METHODS # ========================================================================= def _run_agent(self, agent_name: str) -> AgentResult: """Run a specific agent""" agent = self.agents.get(agent_name) if not agent: logger.warning(f"⚠️ Agent not found: {agent_name}") return AgentResult( status=AgentStatus.SUCCESS, # Skip gracefully agent_name=agent_name, phase=self.current_phase ) self._check_cancelled() result = agent.run() # Log execution self.execution_log.append({ "agent": agent_name, "status": result.status.value, "duration": result.duration_seconds, "timestamp": datetime.now().isoformat() }) # Process any messages for msg in result.messages: self.message_queue.append(msg) return result def _check_cancelled(self): """Check if pipeline was cancelled""" if self.check_cancellation: self.check_cancellation() def _report_progress(self, stage: str, progress: float): """Report progress to callback""" logger.info(f"📍 {stage} ({progress*100:.0f}%)") if self.on_progress: self.on_progress(stage, progress) def _build_result(self, phase_result: Dict, approved: bool) -> Dict[str, Any]: """Build final result dictionary""" duration = (datetime.now() - self.start_time).total_seconds() if self.start_time else 0 return { "success": True, "approved": approved, "pipeline_id": self.memory.pipeline_id, "status": self.status.value, "phase": phase_result.get("phase", "unknown"), "score": phase_result.get("score", 0), "metrics": phase_result.get("metrics", {}), "duration_seconds": duration, "iterations": self.iteration, "model": self.memory.best_model, "execution_log": self.execution_log } def _build_failure_result(self, errors: List[str]) -> Dict[str, Any]: """Build failure result dictionary""" duration = (datetime.now() - self.start_time).total_seconds() if self.start_time else 0 return { "success": False, "approved": False, "pipeline_id": self.memory.pipeline_id, "status": PipelineStatus.FAILED.value, "errors": errors, "duration_seconds": duration, "iterations": self.iteration, "execution_log": self.execution_log }