runtime_environment.py

"""
Runtime Environment - Layer 3: Execution environment for agents
"""

import asyncio
import threading
import time
from typing import Dict, List, Any, Optional, Callable
from datetime import datetime, timedelta
from dataclasses import dataclass
from enum import Enum
import psutil
import os


class AgentStatus(Enum):
    """Agent runtime status"""
    INITIALIZING = "initializing"
    READY = "ready"
    PROCESSING = "processing"
    IDLE = "idle"
    ERROR = "error"
    TERMINATED = "terminated"


@dataclass
class RuntimeMetrics:
    """Runtime performance metrics"""
    timestamp: datetime
    cpu_percent: float
    memory_mb: float
    active_agents: int
    requests_processed: int
    avg_response_time: float
    error_rate: float


class AgentRuntime:
    """Runtime environment for an individual agent"""
    
    def __init__(self, agent_id: str, agent_instance: Any):
        self.agent_id = agent_id
        self.agent = agent_instance
        self.status = AgentStatus.INITIALIZING
        self.metrics = {
            'requests_processed': 0,
            'total_processing_time': 0,
            'errors': 0,
            'last_active': None
        }
        self.task_queue = asyncio.Queue()
        self.processing_lock = threading.Lock()
        self._stop_event = threading.Event()
        
    async def initialize(self):
        """Initialize agent runtime"""
        try:
            # Initialize agent if it has an initialize method
            if hasattr(self.agent, 'initialize'):
                await self.agent.initialize()
            
            self.status = AgentStatus.READY
            self.metrics['last_active'] = datetime.now()
            print(f"✅ Agent {self.agent_id} initialized successfully")
        except Exception as e:
            self.status = AgentStatus.ERROR
            print(f"❌ Agent {self.agent_id} initialization failed: {e}")
    
    async def process_request(self, request: Dict[str, Any]) -> Dict[str, Any]:
        """Process a single request"""
        self.status = AgentStatus.PROCESSING
        start_time = time.time()
        
        try:
            # Call agent's process method
            if hasattr(self.agent, 'process'):
                result = await self.agent.process(**request)
            else:
                result = {'error': 'Agent has no process method', 'response': ''}
            
            # Update metrics
            processing_time = time.time() - start_time
            with self.processing_lock:
                self.metrics['requests_processed'] += 1
                self.metrics['total_processing_time'] += processing_time
                self.metrics['last_active'] = datetime.now()
            
            self.status = AgentStatus.READY
            return result
            
        except Exception as e:
            with self.processing_lock:
                self.metrics['errors'] += 1
            
            self.status = AgentStatus.ERROR
            return {
                'error': str(e),
                'response': f"Agent {self.agent_id} encountered an error",
                'agent_id': self.agent_id
            }
    
    def get_metrics(self) -> Dict[str, Any]:
        """Get agent runtime metrics"""
        with self.processing_lock:
            avg_time = 0
            if self.metrics['requests_processed'] > 0:
                avg_time = self.metrics['total_processing_time'] / self.metrics['requests_processed']
            
            return {
                'agent_id': self.agent_id,
                'status': self.status.value,
                'requests_processed': self.metrics['requests_processed'],
                'avg_processing_time': avg_time,
                'error_rate': self.metrics['errors'] / max(self.metrics['requests_processed'], 1),
                'last_active': self.metrics['last_active'].isoformat() if self.metrics['last_active'] else None,
                'uptime': (datetime.now() - self.metrics['last_active']).total_seconds() if self.metrics['last_active'] else 0
            }
    
    async def cleanup(self):
        """Cleanup agent resources"""
        self.status = AgentStatus.TERMINATED
        if hasattr(self.agent, 'cleanup'):
            await self.agent.cleanup()
        print(f"🛑 Agent {self.agent_id} terminated")


class RuntimeEnvironment:
    """Execution environment for multiple agents"""
    
    def __init__(self, max_agents: int = 10):
        self.agents: Dict[str, AgentRuntime] = {}
        self.max_agents = max_agents
        self.metrics_history: List[RuntimeMetrics] = []
        self.request_log: List[Dict[str, Any]] = []
        self._monitoring_thread: Optional[threading.Thread] = None
        self._stop_monitoring = threading.Event()
        
        # Performance thresholds
        self.thresholds = {
            'max_cpu_percent': 80.0,
            'max_memory_mb': 4096,
            'max_response_time': 5.0,  # seconds
            'max_error_rate': 0.1  # 10%
        }
        
    def register_agent(self, agent_id: str, agent_instance: Any) -> bool:
        """Register a new agent"""
        if len(self.agents) >= self.max_agents:
            print(f"⚠️ Maximum agent limit reached ({self.max_agents})")
            return False
        
        if agent_id in self.agents:
            print(f"⚠️ Agent {agent_id} already registered")
            return False
        
        agent_runtime = AgentRuntime(agent_id, agent_instance)
        self.agents[agent_id] = agent_runtime
        
        # Initialize agent asynchronously
        asyncio.create_task(agent_runtime.initialize())
        
        print(f"✅ Agent {agent_id} registered successfully")
        return True
    
    async def execute_request(self, 
                            agent_id: str, 
                            request: Dict[str, Any]) -> Dict[str, Any]:
        """Execute a request through a specific agent"""
        if agent_id not in self.agents:
            return {'error': f'Agent {agent_id} not found', 'response': ''}
        
        agent = self.agents[agent_id]
        
        # Check if agent is ready
        if agent.status not in [AgentStatus.READY, AgentStatus.IDLE]:
            return {
                'error': f'Agent {agent_id} is not ready (status: {agent.status.value})',
                'response': ''
            }
        
        # Log request
        request_log_entry = {
            'agent_id': agent_id,
            'request': request,
            'timestamp': datetime.now().isoformat(),
            'status': 'processing'
        }
        self.request_log.append(request_log_entry)
        
        # Execute request
        result = await agent.process_request(request)
        
        # Update log
        request_log_entry['status'] = 'completed'
        request_log_entry['processing_time'] = time.time() - datetime.fromisoformat(request_log_entry['timestamp']).timestamp()
        request_log_entry['result'] = {'success': 'error' not in result}
        
        return result
    
    async def execute_parallel(self, 
                             requests: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        """Execute multiple requests in parallel"""
        tasks = []
        
        for req in requests:
            agent_id = req.get('agent_id')
            if agent_id and agent_id in self.agents:
                task = self.execute_request(agent_id, req)
                tasks.append(task)
        
        # Execute in parallel
        results = await asyncio.gather(*tasks, return_exceptions=True)
        
        # Convert exceptions to error responses
        final_results = []
        for result in results:
            if isinstance(result, Exception):
                final_results.append({
                    'error': str(result),
                    'response': 'Parallel execution failed'
                })
            else:
                final_results.append(result)
        
        return final_results
    
    def start_monitoring(self, interval_seconds: int = 30):
        """Start monitoring thread"""
        if self._monitoring_thread and self._monitoring_thread.is_alive():
            print("⚠️ Monitoring already running")
            return
        
        self._stop_monitoring.clear()
        self._monitoring_thread = threading.Thread(
            target=self._monitoring_loop,
            args=(interval_seconds,),
            daemon=True
        )
        self._monitoring_thread.start()
        print(f"📊 Started monitoring with {interval_seconds}s interval")
    
    def _monitoring_loop(self, interval_seconds: int):
        """Monitoring loop running in separate thread"""
        while not self._stop_monitoring.is_set():
            try:
                metrics = self._collect_metrics()
                self.metrics_history.append(metrics)
                
                # Check thresholds
                self._check_thresholds(metrics)
                
                # Keep only last 24 hours of metrics
                cutoff = datetime.now() - timedelta(hours=24)
                self.metrics_history = [
                    m for m in self.metrics_history 
                    if m.timestamp > cutoff
                ]
                
                time.sleep(interval_seconds)
                
            except Exception as e:
                print(f"❌ Monitoring error: {e}")
                time.sleep(interval_seconds)
    
    def _collect_metrics(self) -> RuntimeMetrics:
        """Collect runtime metrics"""
        # System metrics
        cpu_percent = psutil.cpu_percent(interval=0.1)
        memory_info = psutil.Process(os.getpid()).memory_info()
        memory_mb = memory_info.rss / 1024 / 1024
        
        # Agent metrics
        active_agents = sum(1 for agent in self.agents.values() 
                          if agent.status == AgentStatus.PROCESSING)
        
        # Request metrics
        recent_requests = [
            r for r in self.request_log[-100:]  # Last 100 requests
            if datetime.fromisoformat(r['timestamp']) > datetime.now() - timedelta(minutes=5)
        ]
        
        requests_processed = len(recent_requests)
        avg_response_time = 0
        error_rate = 0
        
        if recent_requests:
            processing_times = [
                r.get('processing_time', 0) 
                for r in recent_requests 
                if 'processing_time' in r
            ]
            if processing_times:
                avg_response_time = sum(processing_times) / len(processing_times)
            
            errors = sum(1 for r in recent_requests 
                        if r.get('result', {}).get('success') is False)
            error_rate = errors / len(recent_requests) if recent_requests else 0
        
        return RuntimeMetrics(
            timestamp=datetime.now(),
            cpu_percent=cpu_percent,
            memory_mb=memory_mb,
            active_agents=active_agents,
            requests_processed=requests_processed,
            avg_response_time=avg_response_time,
            error_rate=error_rate
        )
    
    def _check_thresholds(self, metrics: RuntimeMetrics):
        """Check metrics against thresholds and trigger alerts"""
        alerts = []
        
        if metrics.cpu_percent > self.thresholds['max_cpu_percent']:
            alerts.append(f"⚠️ High CPU usage: {metrics.cpu_percent:.1f}%")
        
        if metrics.memory_mb > self.thresholds['max_memory_mb']:
            alerts.append(f"⚠️ High memory usage: {metrics.memory_mb:.1f} MB")
        
        if metrics.avg_response_time > self.thresholds['max_response_time']:
            alerts.append(f"⚠️ Slow response time: {metrics.avg_response_time:.2f}s")
        
        if metrics.error_rate > self.thresholds['max_error_rate']:
            alerts.append(f"⚠️ High error rate: {metrics.error_rate:.1%}")
        
        if alerts:
            print("\n".join(alerts))
            # In production, send alerts to monitoring system
    
    def stop_monitoring(self):
        """Stop monitoring thread"""
        if self._monitoring_thread:
            self._stop_monitoring.set()
            self._monitoring_thread.join(timeout=5)
            print("🛑 Monitoring stopped")
    
    def get_agent_statuses(self) -> Dict[str, Dict[str, Any]]:
        """Get status of all agents"""
        statuses = {}
        for agent_id, agent_runtime in self.agents.items():
            statuses[agent_id] = agent_runtime.get_metrics()
        return statuses
    
    def get_system_metrics(self, time_range: str = '1h') -> Dict[str, Any]:
        """Get system metrics for a time range"""
        if not self.metrics_history:
            return {
                'cpu_percent': 0,
                'memory_mb': 0,
                'active_agents': 0,
                'requests_per_minute': 0,
                'avg_response_time': 0,
                'error_rate': 0
            }
        
        # Calculate time cutoff
        if time_range == '1h':
            cutoff = datetime.now() - timedelta(hours=1)
        elif time_range == '24h':
            cutoff = datetime.now() - timedelta(hours=24)
        else:
            cutoff = datetime.now() - timedelta(minutes=5)
        
        # Filter metrics
        recent_metrics = [m for m in self.metrics_history if m.timestamp > cutoff]
        
        if not recent_metrics:
            # Use latest metric
            recent_metrics = [self.metrics_history[-1]] if self.metrics_history else []
        
        # Calculate averages
        cpu_avg = sum(m.cpu_percent for m in recent_metrics) / len(recent_metrics)
        memory_avg = sum(m.memory_mb for m in recent_metrics) / len(recent_metrics)
        active_avg = sum(m.active_agents for m in recent_metrics) / len(recent_metrics)
        response_avg = sum(m.avg_response_time for m in recent_metrics) / len(recent_metrics)
        error_avg = sum(m.error_rate for m in recent_metrics) / len(recent_metrics)
        
        # Calculate requests per minute
        total_requests = sum(m.requests_processed for m in recent_metrics)
        time_span_hours = (datetime.now() - cutoff).total_seconds() / 3600
        requests_per_minute = (total_requests / time_span_hours) / 60 if time_span_hours > 0 else 0
        
        return {
            'cpu_percent': cpu_avg,
            'memory_mb': memory_avg,
            'active_agents': active_avg,
            'requests_per_minute': requests_per_minute,
            'avg_response_time': response_avg,
            'error_rate': error_avg
        }
    
    def cleanup(self):
        """Cleanup all agents"""
        self.stop_monitoring()
        
        # Cleanup all agents
        cleanup_tasks = []
        for agent_id, agent_runtime in self.agents.items():
            cleanup_tasks.append(agent_runtime.cleanup())
        
        # Wait for cleanup
        if cleanup_tasks:
            asyncio.run(asyncio.gather(*cleanup_tasks))
        
        print("🧹 Runtime environment cleaned up")
    
    def get_performance_report(self) -> Dict[str, Any]:
        """Generate performance report"""
        if not self.metrics_history:
            return {'error': 'No metrics available'}
        
        # Get latest metrics
        latest = self.metrics_history[-1] if self.metrics_history else None
        
        # Calculate trends (last hour vs previous hour)
        one_hour_ago = datetime.now() - timedelta(hours=1)
        two_hours_ago = datetime.now() - timedelta(hours=2)
        
        recent_hour = [m for m in self.metrics_history if m.timestamp > one_hour_ago]
        previous_hour = [m for m in self.metrics_history 
                        if two_hours_ago < m.timestamp <= one_hour_ago]
        
        trend = 'stable'
        if recent_hour and previous_hour:
            recent_avg = sum(m.requests_processed for m in recent_hour) / len(recent_hour)
            previous_avg = sum(m.requests_processed for m in previous_hour) / len(previous_hour)
            
            if recent_avg > previous_avg * 1.2:
                trend = 'improving'
            elif recent_avg < previous_avg * 0.8:
                trend = 'declining'
        
        # Agent performance
        agent_performance = {}
        for agent_id, agent_runtime in self.agents.items():
            agent_metrics = agent_runtime.get_metrics()
            agent_performance[agent_id] = {
                'status': agent_metrics['status'],
                'requests_processed': agent_metrics['requests_processed'],
                'avg_processing_time': agent_metrics['avg_processing_time'],
                'availability': 'high' if agent_metrics['error_rate'] < 0.05 else 'medium' if agent_metrics['error_rate'] < 0.2 else 'low'
            }
        
        return {
            'timestamp': datetime.now().isoformat(),
            'system_metrics': latest.__dict__ if latest else {},
            'trend': trend,
            'agent_performance': agent_performance,
            'alerts': self._generate_alerts() if latest else [],
            'recommendations': self._generate_recommendations()
        }
    
    def _generate_alerts(self) -> List[str]:
        """Generate alerts based on current state"""
        alerts = []
        system_metrics = self.get_system_metrics('5m')
        
        if system_metrics['cpu_percent'] > 70:
            alerts.append("Consider scaling down compute-intensive agents")
        
        if system_metrics['memory_mb'] > 3000:
            alerts.append("High memory usage - check for memory leaks")
        
        if system_metrics['error_rate'] > 0.1:
            alerts.append("High error rate - investigate agent failures")
        
        return alerts
    
    def _generate_recommendations(self) -> List[str]:
        """Generate optimization recommendations"""
        recommendations = []
        agent_statuses = self.get_agent_statuses()
        
        # Find underutilized agents
        for agent_id, status in agent_statuses.items():
            if status['requests_processed'] < 10 and status['uptime'] > 3600:  # < 10 requests in >1 hour
                recommendations.append(f"Agent {agent_id} is underutilized - consider consolidating")
        
        # Check for slow agents
        slow_agents = [
            agent_id for agent_id, status in agent_statuses.items()
            if status['avg_processing_time'] > 2.0  # >2 seconds average
        ]
        if slow_agents:
            recommendations.append(f"Agents {', '.join(slow_agents)} are slow - optimize processing logic")
        
        # Load balancing suggestions
        request_counts = [status['requests_processed'] for status in agent_statuses.values()]
        if request_counts:
            max_requests = max(request_counts)
            min_requests = min(request_counts)
            if max_requests > min_requests * 3:  # Significant imbalance
                recommendations.append("Request load is imbalanced - implement load balancing")
        
        return recommendations