Fix: [Enhancement] Agent Fallback and Recovery Mechanisms (Auto-Generated)

agent_fallback.py

@@ -0,0 +1,325 @@
+"""Agent Fallback and Recovery Mechanisms (#225).
+
+Provides robust resilience patterns for agentic workflows:
+- Circuit Breaker pattern to prevent cascading failures
+- Exponential backoff retry logic with jitter
+- Health check and monitoring for agent types
+- Alternative agent routing and graceful degradation
+- Partial result preservation and caching
+- Deterministic / rule-based fallbacks
+- Failure analytics and telemetry logging
+"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import random
+import time
+from collections.abc import Awaitable, Callable
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+
+class CircuitState(str, Enum):
+    CLOSED = "closed"
+    OPEN = "open"
+    HALF_OPEN = "half_open"
+
+
+class AgentHealth(str, Enum):
+    HEALTHY = "healthy"
+    DEGRADED = "degraded"
+    UNHEALTHY = "unhealthy"
+
+
+@dataclass
+class CircuitBreakerConfig:
+    failure_threshold: int = 3
+    recovery_timeout_seconds: float = 30.0
+    half_open_max_calls: int = 2
+
+
+@dataclass
+class RetryConfig:
+    max_retries: int = 3
+    initial_delay: float = 0.1
+    max_delay: float = 5.0
+    backoff_factor: float = 2.0
+    jitter: bool = True
+
+
+class CircuitBreaker:
+    """Circuit breaker to prevent cascading failures."""
+
+    def __init__(self, name: str, config: CircuitBreakerConfig | None = None):
+        self.name = name
+        self.config = config or CircuitBreakerConfig()
+        self.state = CircuitState.CLOSED
+        self.failure_count = 0
+        self.success_count = 0
+        self.last_failure_time = 0.0
+        self.half_open_calls = 0
+
+    def record_success(self) -> None:
+        if self.state == CircuitState.HALF_OPEN:
+            self.success_count += 1
+            if self.success_count >= self.config.half_open_max_calls:
+                self.state = CircuitState.CLOSED
+                self.failure_count = 0
+                self.success_count = 0
+                logger.info(f"Circuit breaker '{self.name}' recovered and is now CLOSED.")
+        elif self.state == CircuitState.CLOSED:
+            self.failure_count = 0
+
+    def record_failure(self) -> None:
+        self.failure_count += 1
+        self.last_failure_time = time.time()
+        if self.state == CircuitState.HALF_OPEN:
+            self.state = CircuitState.OPEN
+            self.success_count = 0
+            logger.warning(f"Circuit breaker '{self.name}' failed in HALF_OPEN, returning to OPEN.")
+        elif self.state == CircuitState.CLOSED and self.failure_count >= self.config.failure_threshold:
+            self.state = CircuitState.OPEN
+            logger.warning(f"Circuit breaker '{self.name}' threshold reached ({self.failure_count}), now OPEN.")
+
+    def allow_request(self) -> bool:
+        if self.state == CircuitState.CLOSED:
+            return True
+        if self.state == CircuitState.OPEN:
+            now = time.time()
+            if now - self.last_failure_time >= self.config.recovery_timeout_seconds:
+                self.state = CircuitState.HALF_OPEN
+                self.half_open_calls = 0
+                self.success_count = 0
+                logger.info(f"Circuit breaker '{self.name}' recovery timeout elapsed, testing in HALF_OPEN.")
+                return True
+            return False
+        if self.state == CircuitState.HALF_OPEN:
+            return True
+        return True
+
+
+@dataclass
+class AgentExecutionResult:
+    success: bool
+    data: Any | None = None
+    error: str | None = None
+    partial_result: Any | None = None
+    stages_executed: list[str] = field(default_factory=list)
+    fallback_used: bool = False
+    duration_ms: float = 0.0
+
+
+class AgentRegistry:
+    """Registry for agents, fallback routes, health status, and failure metrics."""
+
+    def __init__(self) -> None:
+        self._agents: dict[str, Callable[..., Awaitable[Any]]] = {}
+        self._fallbacks: dict[str, str] = {}  # primary -> backup agent name
+        self._deterministic_fallbacks: dict[str, Callable[..., Any]] = {}
+        self._circuit_breakers: dict[str, CircuitBreaker] = {}
+        self._health_statuses: dict[str, AgentHealth] = {}
+        self._failure_counts: dict[str, int] = {}
+        self._success_counts: dict[str, int] = {}
+        self._cached_results: dict[str, tuple[float, Any]] = {}
+        self.cache_ttl: float = 300.0
+
+    def register_agent(
+        self,
+        name: str,
+        handler: Callable[..., Awaitable[Any]],
+        fallback_agent: str | None = None,
+        deterministic_fallback: Callable[..., Any] | None = None,
+        circuit_config: CircuitBreakerConfig | None = None,
+    ) -> None:
+        self._agents[name] = handler
+        if fallback_agent:
+            self._fallbacks[name] = fallback_agent
+        if deterministic_fallback:
+            self._deterministic_fallbacks[name] = deterministic_fallback
+        self._circuit_breakers[name] = CircuitBreaker(name, circuit_config)
+        self._health_statuses[name] = AgentHealth.HEALTHY
+        self._failure_counts[name] = 0
+        self._success_counts[name] = 0
+
+    def get_health(self, name: str) -> AgentHealth:
+        cb = self._circuit_breakers.get(name)
+        if cb and cb.state == CircuitState.OPEN:
+            return AgentHealth.UNHEALTHY
+        fails = self._failure_counts.get(name, 0)
+        succs = self._success_counts.get(name, 1)
+        if fails / (fails + succs) > 0.3 and fails >= 2:
+            return AgentHealth.DEGRADED
+        return AgentHealth.HEALTHY
+
+    def get_all_health(self) -> dict[str, str]:
+        return {name: self.get_health(name).value for name in self._agents}
+
+    def record_failure(self, name: str) -> None:
+        self._failure_counts[name] = self._failure_counts.get(name, 0) + 1
+        if name in self._circuit_breakers:
+            self._circuit_breakers[name].record_failure()
+        logger.error(f"Agent '{name}' recorded failure. Total failures: {self._failure_counts[name]}")
+
+    def record_success(self, name: str) -> None:
+        self._success_counts[name] = self._success_counts.get(name, 0) + 1
+        if name in self._circuit_breakers:
+            self._circuit_breakers[name].record_success()
+
+    def cache_result(self, key: str, data: Any) -> None:
+        self._cached_results[key] = (time.time(), data)
+
+    def get_cached_result(self, key: str) -> Any | None:
+        if key in self._cached_results:
+            ts, data = self._cached_results[key]
+            if time.time() - ts <= self.cache_ttl:
+                return data
+        return None
+
+    async def execute_with_resilience(
+        self,
+        agent_name: str,
+        payload: Any,
+        retry_config: RetryConfig | None = None,
+        cache_key: str | None = None,
+        partial_preserver: Callable[[Any, Exception], Any] | None = None,
+    ) -> AgentExecutionResult:
+        start_time = time.time()
+        stages = []
+        retry_cfg = retry_config or RetryConfig()
+
+        if cache_key:
+            cached = self.get_cached_result(cache_key)
+            if cached is not None:
+                stages.append("cache_hit")
+                return AgentExecutionResult(
+                    success=True,
+                    data=cached,
+                    stages_executed=stages,
+                    duration_ms=(time.time() - start_time) * 1000,
+                )
+
+        cb = self._circuit_breakers.get(agent_name)
+        if cb and not cb.allow_request():
+            stages.append("circuit_breaker_open")
+            logger.warning(f"Circuit breaker for '{agent_name}' is OPEN. Attempting fallback routing.")
+            return await self._execute_fallback_chain(
+                agent_name, payload, start_time, stages, Exception("Circuit breaker is open"), partial_preserver
+            )
+
+        handler = self._agents.get(agent_name)
+        if not handler:
+            return AgentExecutionResult(
+                success=False,
+                error=f"Agent '{agent_name}' not found",
+                stages_executed=stages,
+                duration_ms=(time.time() - start_time) * 1000,
+            )
+
+        last_exception = None
+        for attempt in range(retry_cfg.max_retries + 1):
+            try:
+                stages.append(f"attempt_{attempt}")
+                res = await handler(payload)
+                self.record_success(agent_name)
+                if cache_key:
+                    self.cache_result(cache_key, res)
+                stages.append("success")
+                return AgentExecutionResult(
+                    success=True,
+                    data=res,
+                    stages_executed=stages,
+                    duration_ms=(time.time() - start_time) * 1000,
+                )
+            except Exception as e:
+                last_exception = e
+                logger.warning(f"Agent '{agent_name}' attempt {attempt} failed: {e}")
+                if attempt < retry_cfg.max_retries:
+                    delay = retry_cfg.initial_delay * (retry_cfg.backoff_factor ** attempt)
+                    delay = min(delay, retry_cfg.max_delay)
+                    if retry_cfg.jitter:
+                        delay *= 0.5 + random.random()
+                    await asyncio.sleep(delay)
+
+        self.record_failure(agent_name)
+        return await self._execute_fallback_chain(
+            agent_name, payload, start_time, stages, last_exception, partial_preserver
+        )
+
+    async def _execute_fallback_chain(
+        self,
+        agent_name: str,
+        payload: Any,
+        start_time: float,
+        stages: list[str],
+        last_exception: Exception | None,
+        partial_preserver: Callable[[Any, Exception], Any] | None,
+    ) -> AgentExecutionResult:
+        partial = None
+        if partial_preserver and last_exception:
+            try:
+                partial = partial_preserver(payload, last_exception)
+                stages.append("partial_result_preserved")
+            except Exception as pe:
+                logger.error(f"Partial preserver failed: {pe}")
+
+        # 1. Try alternative agent route
+        alt_agent = self._fallbacks.get(agent_name)
+        if alt_agent and alt_agent in self._agents:
+            stages.append(f"routed_to_alternative_{alt_agent}")
+            logger.info(f"Routing failed agent '{agent_name}' to alternative agent '{alt_agent}'.")
+            try:
+                alt_handler = self._agents[alt_agent]
+                res = await alt_handler(payload)
+                self.record_success(alt_agent)
+                return AgentExecutionResult(
+                    success=True,
+                    data=res,
+                    partial_result=partial,
+                    stages_executed=stages,
+                    fallback_used=True,
+                    duration_ms=(time.time() - start_time) * 1000,
+                )
+            except Exception as ae:
+                logger.error(f"Alternative agent '{alt_agent}' also failed: {ae}")
+
+        # 2. Try deterministic / rule-based fallback
+        det_fallback = self._deterministic_fallbacks.get(agent_name)
+        if det_fallback:
+            stages.append("deterministic_fallback_used")
+            logger.info(f"Falling back to deterministic method for agent '{agent_name}'.")
+            try:
+                if asyncio.iscoroutinefunction(det_fallback):
+                    res = await det_fallback(payload)
+                else:
+                    res = det_fallback(payload)
+                return AgentExecutionResult(
+                    success=True,
+                    data=res,
+                    partial_result=partial,
+                    stages_executed=stages,
+                    fallback_used=True,
+                    duration_ms=(time.time() - start_time) * 1000,
+                )
+            except Exception as de:
+                logger.error(f"Deterministic fallback failed: {de}")
+
+        # 3. Graceful degradation failure
+        stages.append("graceful_degradation_failed")
+        err_msg = str(last_exception) if last_exception else "Unknown failure"
+        return AgentExecutionResult(
+            success=False,
+            error=err_msg,
+            partial_result=partial,
+            stages_executed=stages,
+            fallback_used=True,
+            duration_ms=(time.time() - start_time) * 1000,
+        )
+
+
+# Global singleton registry instance
+default_agent_registry = AgentRegistry()