Building Production-Ready AI Agents with MCP Architecture

import asyncio
import json
import logging
from datetime import datetime, timedelta
from typing import Dict, List, Any, Optional, Callable
from dataclasses import dataclass, asdict
from enum import Enum
import uuid
import aiohttp
import redis.asyncio as redis
from contextlib import asynccontextmanager

# Agent State Management
class AgentState(Enum):
    IDLE = "idle"
    THINKING = "thinking"
    EXECUTING = "executing"
    ERROR = "error"
    COMPLETED = "completed"

@dataclass
class AgentContext:
    session_id: str
    user_id: str
    conversation_history: List[Dict]
    available_tools: List[str]
    state: AgentState
    created_at: datetime
    last_activity: datetime
    metadata: Dict[str, Any]

@dataclass
class AgentResponse:
    response_id: str
    content: str
    tool_calls: List[Dict]
    state: AgentState
    confidence: float
    execution_time: float
    metadata: Dict[str, Any]

class ProductionAIAgent:
    """
    Production-ready AI agent with enterprise features:
    - Robust error handling and recovery
    - Comprehensive logging and metrics
    - Horizontal scalability
    - Security and rate limiting
    - Context management and persistence
    """
    
    def __init__(self, 
                 agent_id: str,
                 llm_client: Any,
                 mcp_registry: 'MCPToolRegistry',
                 redis_client: redis.Redis,
                 config: Dict[str, Any]):
        
        self.agent_id = agent_id
        self.llm_client = llm_client
        self.mcp_registry = mcp_registry
        self.redis = redis_client
        self.config = config
        
        # Initialize logging with structured format
        self.logger = self._setup_logger()
        
        # Performance and reliability tracking
        self.metrics = AgentMetrics()
        self.circuit_breaker = CircuitBreaker(
            failure_threshold=config.get('circuit_breaker_threshold', 5),
            recovery_timeout=config.get('circuit_breaker_timeout', 60)
        )
        
        # Context and session management
        self.context_manager = ContextManager(redis_client)
        self.rate_limiter = RateLimiter(redis_client)
    
    async def process_request(self, 
                            request: Dict[str, Any], 
                            context: AgentContext) -> AgentResponse:
        """
        Main request processing pipeline with comprehensive error handling
        """
        start_time = datetime.utcnow()
        response_id = str(uuid.uuid4())
        
        try:
            # Rate limiting check
            if not await self.rate_limiter.allow_request(context.user_id):
                raise AgentError("Rate limit exceeded", code="RATE_LIMITED")
            
            # Update agent state
            context.state = AgentState.THINKING
            await self._persist_context(context)
            
            # Validate and sanitize input
            validated_request = await self._validate_request(request, context)
            
            # Generate response using LLM with circuit breaker protection
            with self.circuit_breaker:
                response = await self._generate_response(validated_request, context)
            
            # Execute tool calls if present
            if response.tool_calls:
                context.state = AgentState.EXECUTING
                await self._persist_context(context)
                
                tool_results = await self._execute_tools(response.tool_calls, context)
                response = await self._incorporate_tool_results(response, tool_results, context)
            
            # Finalize response
            context.state = AgentState.COMPLETED
            context.last_activity = datetime.utcnow()
            await self._persist_context(context)
            
            # Record metrics
            execution_time = (datetime.utcnow() - start_time).total_seconds()
            await self.metrics.record_success(execution_time)
            
            self.logger.info(
                "Request processed successfully",
                extra={
                    "response_id": response_id,
                    "session_id": context.session_id,
                    "execution_time": execution_time,
                    "tool_calls": len(response.tool_calls)
                }
            )
            
            return AgentResponse(
                response_id=response_id,
                content=response.content,
                tool_calls=response.tool_calls,
                state=context.state,
                confidence=response.confidence,
                execution_time=execution_time,
                metadata={"agent_id": self.agent_id}
            )
            
        except Exception as e:
            # Comprehensive error handling
            context.state = AgentState.ERROR
            await self._persist_context(context)
            await self.metrics.record_error(type(e).__name__)
            
            self.logger.error(
                "Request processing failed",
                extra={
                    "response_id": response_id,
                    "session_id": context.session_id,
                    "error": str(e),
                    "error_type": type(e).__name__
                },
                exc_info=True
            )
            
            # Return graceful error response
            return self._create_error_response(response_id, e, context)
    
    async def _generate_response(self, request: Dict, context: AgentContext) -> Any:
        """
        Generate LLM response with retry logic and timeout handling
        """
        max_retries = self.config.get('max_retries', 3)
        retry_delay = self.config.get('retry_delay', 1.0)
        
        for attempt in range(max_retries):
            try:
                # Prepare context for LLM
                llm_context = self._prepare_llm_context(request, context)
                
                # Call LLM with timeout
                async with asyncio.timeout(self.config.get('llm_timeout', 30)):
                    response = await self.llm_client.generate(
                        messages=llm_context['messages'],
                        tools=llm_context['available_tools'],
                        max_tokens=self.config.get('max_tokens', 2000)
                    )
                
                return response
                
            except asyncio.TimeoutError:
                if attempt < max_retries - 1:
                    await asyncio.sleep(retry_delay * (2 ** attempt))  # Exponential backoff
                    continue
                raise AgentError("LLM request timeout", code="LLM_TIMEOUT")
                
            except Exception as e:
                if attempt < max_retries - 1:
                    await asyncio.sleep(retry_delay * (2 ** attempt))
                    continue
                raise AgentError(f"LLM generation failed: {str(e)}", code="LLM_ERROR")
        
    async def _execute_tools(self, tool_calls: List[Dict], context: AgentContext) -> List[Dict]:
        """
        Execute MCP tools with parallel processing and error isolation
        """
        results = []
        semaphore = asyncio.Semaphore(self.config.get('max_concurrent_tools', 5))
        
        async def execute_single_tool(tool_call: Dict) -> Dict:
            async with semaphore:
                try:
                    tool_name = tool_call['name']
                    parameters = tool_call['parameters']
                    
                    # Validate tool availability and permissions
                    if not await self._validate_tool_access(tool_name, context):
                        return {
                            "tool_name": tool_name,
                            "success": False,
                            "error": "Tool access denied",
                            "result": None
                        }
                    
                    # Execute tool through MCP registry
                    async with asyncio.timeout(self.config.get('tool_timeout', 10)):
                        result = await self.mcp_registry.execute_tool(
                            tool_name, parameters, context.user_id
                        )
                    
                    return {
                        "tool_name": tool_name,
                        "success": True,
                        "error": None,
                        "result": result
                    }
                    
                except Exception as e:
                    self.logger.warning(
                        f"Tool execution failed: {tool_name}",
                        extra={"error": str(e), "session_id": context.session_id}
                    )
                    
                    return {
                        "tool_name": tool_name,
                        "success": False,
                        "error": str(e),
                        "result": None
                    }
        
        # Execute tools concurrently with error isolation
        tasks = [execute_single_tool(tool_call) for tool_call in tool_calls]
        results = await asyncio.gather(*tasks, return_exceptions=True)
        
        # Handle any task exceptions
        final_results = []
        for i, result in enumerate(results):
            if isinstance(result, Exception):
                final_results.append({
                    "tool_name": tool_calls[i].get('name', 'unknown'),
                    "success": False,
                    "error": str(result),
                    "result": None
                })
            else:
                final_results.append(result)
        
        return final_results
    
    async def _persist_context(self, context: AgentContext):
        """
        Persist agent context to Redis for stateless scalability
        """
        try:
            context_key = f"agent_context:{context.session_id}"
            context_data = json.dumps(asdict(context), default=str)
            
            await self.redis.setex(
                context_key, 
                self.config.get('context_ttl', 3600),  # 1 hour default
                context_data
            )
            
        except Exception as e:
            self.logger.error(f"Failed to persist context: {str(e)}")
            # Don't fail the request for context persistence issues
    
    def _setup_logger(self) -> logging.Logger:
        """
        Set up structured logging for production observability
        """
        logger = logging.getLogger(f"ai_agent.{self.agent_id}")
        logger.setLevel(self.config.get('log_level', logging.INFO))
        
        if not logger.handlers:
            handler = logging.StreamHandler()
            formatter = logging.Formatter(
                '%(asctime)s - %(name)s - %(levelname)s - %(message)s - %(pathname)s:%(lineno)d'
            )
            handler.setFormatter(formatter)
            logger.addHandler(handler)
        
        return logger

class MCPToolRegistry:
    """
    Registry for managing MCP tools with caching and health checks
    """
    
    def __init__(self, redis_client: redis.Redis):
        self.redis = redis_client
        self.tool_cache = {}
        self.health_checks = {}
    
    async def execute_tool(self, tool_name: str, parameters: Dict, user_id: str) -> Any:
        """
        Execute MCP tool with validation, caching, and monitoring
        """
        # Implementation would connect to actual MCP servers
        # This is a simplified example
        
        tool_config = await self._get_tool_config(tool_name)
        if not tool_config:
            raise ToolNotFoundError(f"Tool {tool_name} not found")
        
        # Validate parameters
        validated_params = await self._validate_tool_parameters(
            tool_name, parameters, tool_config
        )
        
        # Execute with monitoring
        start_time = datetime.utcnow()
        try:
            result = await self._execute_mcp_tool(
                tool_config['server_id'], 
                tool_name, 
                validated_params
            )
            
            execution_time = (datetime.utcnow() - start_time).total_seconds()
            await self._record_tool_metrics(tool_name, execution_time, True)
            
            return result
            
        except Exception as e:
            execution_time = (datetime.utcnow() - start_time).total_seconds()
            await self._record_tool_metrics(tool_name, execution_time, False)
            raise

# Supporting Classes for Production Features

class CircuitBreaker:
    """
    Circuit breaker pattern for fault tolerance
    """
    def __init__(self, failure_threshold: int, recovery_timeout: int):
        self.failure_threshold = failure_threshold
        self.recovery_timeout = recovery_timeout
        self.failure_count = 0
        self.last_failure_time = None
        self.state = "closed"  # closed, open, half-open
    
    def __enter__(self):
        if self.state == "open":
            if datetime.utcnow() - self.last_failure_time > timedelta(seconds=self.recovery_timeout):
                self.state = "half-open"
            else:
                raise CircuitBreakerOpenError("Circuit breaker is open")
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        if exc_type is None:
            # Success
            if self.state == "half-open":
                self.state = "closed"
                self.failure_count = 0
        else:
            # Failure
            self.failure_count += 1
            self.last_failure_time = datetime.utcnow()
            
            if self.failure_count >= self.failure_threshold:
                self.state = "open"

class AgentMetrics:
    """
    Metrics collection for agent performance monitoring
    """
    def __init__(self):
        self.request_count = 0
        self.error_count = 0
        self.total_execution_time = 0.0
        self.error_types = {}
    
    async def record_success(self, execution_time: float):
        self.request_count += 1
        self.total_execution_time += execution_time
    
    async def record_error(self, error_type: str):
        self.error_count += 1
        self.error_types[error_type] = self.error_types.get(error_type, 0) + 1
    
    def get_metrics(self) -> Dict[str, Any]:
        return {
            "request_count": self.request_count,
            "error_count": self.error_count,
            "error_rate": self.error_count / max(self.request_count, 1),
            "avg_execution_time": self.total_execution_time / max(self.request_count, 1),
            "error_breakdown": self.error_types
        }

# Custom Exceptions
class AgentError(Exception):
    def __init__(self, message: str, code: str = None):
        self.message = message
        self.code = code
        super().__init__(self.message)

class CircuitBreakerOpenError(AgentError):
    pass

class ToolNotFoundError(AgentError):
    pass

class ErrorRecoveryManager:
    """
    Comprehensive error recovery strategies for production AI agents
    """
    
    def __init__(self, config: Dict[str, Any]):
        self.config = config
        self.recovery_strategies = {
            "LLM_TIMEOUT": self._handle_llm_timeout,
            "TOOL_FAILURE": self._handle_tool_failure,
            "RATE_LIMITED": self._handle_rate_limit,
            "CONTEXT_LOST": self._handle_context_loss,
            "RESOURCE_EXHAUSTED": self._handle_resource_exhaustion
        }
    
    async def recover_from_error(self, error: AgentError, context: AgentContext) -> AgentResponse:
        """
        Apply appropriate recovery strategy based on error type
        """
        strategy = self.recovery_strategies.get(error.code)
        if strategy:
            return await strategy(error, context)
        else:
            return await self._handle_generic_error(error, context)
    
    async def _handle_llm_timeout(self, error: AgentError, context: AgentContext) -> AgentResponse:
        """
        Recovery strategy for LLM timeout errors
        """
        # Try with reduced context or simpler prompt
        simplified_context = self._simplify_context(context)
        
        try:
            # Retry with shorter context and lower token limit
            response = await self._retry_with_simplified_request(simplified_context)
            
            return AgentResponse(
                response_id=str(uuid.uuid4()),
                content=response.content + "\n\n*Note: Response generated with simplified context due to timeout.*",
                tool_calls=[],
                state=AgentState.COMPLETED,
                confidence=0.7,  # Lower confidence due to simplified context
                execution_time=0.0,
                metadata={"recovery_applied": "simplified_context"}
            )
            
        except Exception:
            # Final fallback: return cached or template response
            return await self._get_fallback_response(context)
    
    async def _handle_tool_failure(self, error: AgentError, context: AgentContext) -> AgentResponse:
        """
        Recovery strategy for tool execution failures
        """
        # Provide helpful error message and suggest alternatives
        fallback_content = f"""I encountered an issue executing the requested tool: {error.message}

Here are some alternatives you can try:
1. Check if the tool parameters are correct
2. Verify you have necessary permissions
3. Try a similar tool if available
4. Contact support if the issue persists

How would you like to proceed?"""
        
        return AgentResponse(
            response_id=str(uuid.uuid4()),
            content=fallback_content,
            tool_calls=[],
            state=AgentState.ERROR,
            confidence=0.8,
            execution_time=0.0,
            metadata={"recovery_applied": "tool_failure_guidance"}
        )
    
    async def _handle_context_loss(self, error: AgentError, context: AgentContext) -> AgentResponse:
        """
        Recovery strategy when agent context is lost
        """
        recovery_content = """I apologize, but I've lost the context of our previous conversation. This can happen due to:

- Session timeout
- Server maintenance
- High system load

To continue effectively, could you please:
1. Briefly remind me what we were discussing
2. Restate your current question or request
3. Provide any relevant context

I'm ready to help once I understand the current situation."""
        
        return AgentResponse(
            response_id=str(uuid.uuid4()),
            content=recovery_content,
            tool_calls=[],
            state=AgentState.IDLE,
            confidence=1.0,
            execution_time=0.0,
            metadata={"recovery_applied": "context_recovery"}
        )

class HealthCheckManager:
    """
    Comprehensive health monitoring for all agent components
    """
    
    def __init__(self, agent: ProductionAIAgent):
        self.agent = agent
        self.health_checks = {
            "llm_connection": self._check_llm_health,
            "mcp_tools": self._check_mcp_tools_health,
            "redis_connection": self._check_redis_health,
            "circuit_breakers": self._check_circuit_breaker_health
        }
    
    async def run_health_checks(self) -> Dict[str, Any]:
        """
        Execute all health checks and return status
        """
        results = {}
        overall_healthy = True
        
        for check_name, check_func in self.health_checks.items():
            try:
                result = await asyncio.wait_for(check_func(), timeout=5.0)
                results[check_name] = {
                    "status": "healthy" if result else "unhealthy",
                    "details": result if isinstance(result, dict) else {}
                }
                if not result:
                    overall_healthy = False
                    
            except Exception as e:
                results[check_name] = {
                    "status": "error",
                    "error": str(e)
                }
                overall_healthy = False
        
        return {
            "overall_status": "healthy" if overall_healthy else "unhealthy",
            "timestamp": datetime.utcnow().isoformat(),
            "components": results,
            "agent_metrics": self.agent.metrics.get_metrics()
        }
    
    async def _check_llm_health(self) -> bool:
        """
        Check LLM service connectivity and response time
        """
        try:
            start_time = datetime.utcnow()
            # Simple health check request
            await self.agent.llm_client.health_check()
            response_time = (datetime.utcnow() - start_time).total_seconds()
            
            return {
                "responsive": True,
                "response_time": response_time,
                "status": "operational"
            }
        except Exception:
            return False

class AgentCacheManager:
    """
    Multi-layer caching strategy for AI agent optimization
    """
    
    def __init__(self, redis_client: redis.Redis, config: Dict[str, Any]):
        self.redis = redis_client
        self.config = config
        self.cache_layers = {
            "response": ResponseCache(redis_client, ttl=300),      # 5 minutes
            "tool_results": ToolResultCache(redis_client, ttl=900), # 15 minutes
            "context": ContextCache(redis_client, ttl=3600),        # 1 hour
            "embeddings": EmbeddingCache(redis_client, ttl=86400)   # 24 hours
        }
    
    async def get_cached_response(self, request_hash: str) -> Optional[AgentResponse]:
        """
        Check for cached response to identical requests
        """
        cache_key = f"agent_response:{request_hash}"
        cached_data = await self.redis.get(cache_key)
        
        if cached_data:
            try:
                response_data = json.loads(cached_data)
                return AgentResponse(**response_data)
            except Exception as e:
                # Invalid cache entry, remove it
                await self.redis.delete(cache_key)
        
        return None
    
    async def cache_response(self, request_hash: str, response: AgentResponse):
        """
        Cache successful responses for future identical requests
        """
        # Only cache successful responses
        if response.state == AgentState.COMPLETED and response.confidence > 0.8:
            cache_key = f"agent_response:{request_hash}"
            response_data = json.dumps(asdict(response), default=str)
            
            await self.redis.setex(
                cache_key,
                self.config.get('response_cache_ttl', 300),
                response_data
            )
    
    async def get_cached_tool_result(self, tool_name: str, params_hash: str) -> Optional[Any]:
        """
        Retrieve cached tool execution results
        """
        cache_key = f"tool_result:{tool_name}:{params_hash}"
        cached_result = await self.redis.get(cache_key)
        
        if cached_result:
            return json.loads(cached_result)
        
        return None
    
    def _generate_request_hash(self, request: Dict[str, Any], context: AgentContext) -> str:
        """
        Generate deterministic hash for request caching
        """
        # Include relevant context that affects response
        cache_input = {
            "message": request.get("message", ""),
            "user_id": context.user_id,
            "available_tools": sorted(context.available_tools),
            "timestamp_hour": datetime.utcnow().strftime("%Y-%m-%d-%H")  # Hourly cache
        }
        
        cache_string = json.dumps(cache_input, sort_keys=True)
        return hashlib.sha256(cache_string.encode()).hexdigest()[:16]

class LoadBalancingManager:
    """
    Intelligent load balancing for agent instances
    """
    
    def __init__(self, redis_client: redis.Redis):
        self.redis = redis_client
        self.agent_instances = {}
        self.health_scores = {}
    
    async def get_optimal_agent(self, request_complexity: str) -> str:
        """
        Select the best agent instance based on current load and capabilities
        """
        available_agents = await self._get_healthy_agents()
        
        if not available_agents:
            raise NoAgentsAvailableError("No healthy agent instances available")
        
        # Simple round-robin for now, can be enhanced with:
        # - Least connections
        # - Weighted round-robin based on instance capacity
        # - Consistent hashing for session affinity
        
        best_agent = min(available_agents, key=lambda a: self._get_agent_load(a))
        
        # Update load tracking
        await self._increment_agent_load(best_agent)
        
        return best_agent
    
    async def _get_healthy_agents(self) -> List[str]:
        """
        Get list of healthy agent instances
        """
        agent_keys = await self.redis.keys("agent_health:*")
        healthy_agents = []
        
        for key in agent_keys:
            health_data = await self.redis.get(key)
            if health_data:
                health_info = json.loads(health_data)
                if health_info.get("status") == "healthy":
                    agent_id = key.decode().split(":")[-1]
                    healthy_agents.append(agent_id)
        
        return healthy_agents