diff --git a/agent_fallback.py b/agent_fallback.py new file mode 100644 index 0000000..c221122 --- /dev/null +++ b/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() diff --git a/tests/test_agent_fallback.py b/tests/test_agent_fallback.py new file mode 100644 index 0000000..29c693d --- /dev/null +++ b/tests/test_agent_fallback.py @@ -0,0 +1,184 @@ +"""Tests for Agent Fallback and Recovery Mechanisms (#225).""" + +import asyncio +import pytest + +from agent_fallback import ( + AgentHealth, + AgentRegistry, + CircuitBreakerConfig, + CircuitState, + RetryConfig, +) + + +@pytest.mark.asyncio +async def test_agent_success_first_try(): + registry = AgentRegistry() + + async def success_agent(payload): + return f"Hello {payload['name']}" + + registry.register_agent("primary", success_agent) + + result = await registry.execute_with_resilience("primary", {"name": "DeenBridge"}) + assert result.success is True + assert result.data == "Hello DeenBridge" + assert "success" in result.stages_executed + assert registry.get_health("primary") == AgentHealth.HEALTHY + + +@pytest.mark.asyncio +async def test_agent_retry_recovery(): + registry = AgentRegistry() + attempts = 0 + + async def flaky_agent(payload): + nonlocal attempts + attempts += 1 + if attempts < 3: + raise ConnectionError("Transient timeout") + return "Recovered Success" + + registry.register_agent("flaky", flaky_agent) + + result = await registry.execute_with_resilience( + "flaky", + {}, + retry_config=RetryConfig(max_retries=3, initial_delay=0.01), + ) + assert result.success is True + assert result.data == "Recovered Success" + assert attempts == 3 + + +@pytest.mark.asyncio +async def test_alternative_agent_routing(): + registry = AgentRegistry() + + async def failing_primary(payload): + raise RuntimeError("Primary crashed") + + async def backup_agent(payload): + return "Backup Result" + + registry.register_agent("primary", failing_primary, fallback_agent="backup") + registry.register_agent("backup", backup_agent) + + result = await registry.execute_with_resilience( + "primary", + {}, + retry_config=RetryConfig(max_retries=1, initial_delay=0.01), + ) + assert result.success is True + assert result.data == "Backup Result" + assert result.fallback_used is True + assert any("routed_to_alternative_backup" in s for s in result.stages_executed) + + +@pytest.mark.asyncio +async def test_deterministic_fallback(): + registry = AgentRegistry() + + async def failing_agent(payload): + raise ValueError("Total failure") + + def rule_based_fallback(payload): + return "Deterministic Rule Output" + + registry.register_agent( + "ai_agent", + failing_agent, + deterministic_fallback=rule_based_fallback, + ) + + result = await registry.execute_with_resilience( + "ai_agent", + {}, + retry_config=RetryConfig(max_retries=0), + ) + assert result.success is True + assert result.data == "Deterministic Rule Output" + assert "deterministic_fallback_used" in result.stages_executed + + +@pytest.mark.asyncio +async def test_circuit_breaker_tripping(): + registry = AgentRegistry() + + async def always_fail(payload): + raise TimeoutError("Timeout") + + registry.register_agent( + "cb_agent", + always_fail, + circuit_config=CircuitBreakerConfig(failure_threshold=2, recovery_timeout_seconds=60.0), + deterministic_fallback=lambda p: "Fallback Rule", + ) + + # First failure + r1 = await registry.execute_with_resilience("cb_agent", {}, retry_config=RetryConfig(max_retries=0)) + assert r1.success is True # fallback caught it + + # Second failure -> trips circuit breaker + r2 = await registry.execute_with_resilience("cb_agent", {}, retry_config=RetryConfig(max_retries=0)) + assert r2.success is True + + cb = registry._circuit_breakers["cb_agent"] + assert cb.state == CircuitState.OPEN + assert registry.get_health("cb_agent") == AgentHealth.UNHEALTHY + + # Third call should immediately hit open circuit breaker without invoking always_fail + r3 = await registry.execute_with_resilience("cb_agent", {}, retry_config=RetryConfig(max_retries=0)) + assert r3.success is True + assert "circuit_breaker_open" in r3.stages_executed + + +@pytest.mark.asyncio +async def test_partial_result_preservation(): + registry = AgentRegistry() + + async def failing_with_partial(payload): + raise RuntimeError("Failed mid-stream") + + def preserve_partial(payload, exc): + return {"partial_text": "partial tokens gathered", "error": str(exc)} + + registry.register_agent("stream_agent", failing_with_partial) + + result = await registry.execute_with_resilience( + "stream_agent", + {}, + retry_config=RetryConfig(max_retries=0), + partial_preserver=preserve_partial, + ) + assert result.success is False + assert result.partial_result is not None + assert result.partial_result["partial_text"] == "partial tokens gathered" + assert "partial_result_preserved" in result.stages_executed + + +@pytest.mark.asyncio +async def test_result_caching_redundancy(): + registry = AgentRegistry() + call_count = 0 + + async def cached_agent(payload): + nonlocal call_count + call_count += 1 + return "Expensive Compute Result" + + registry.register_agent("compute", cached_agent) + + # First call + res1 = await registry.execute_with_resilience("compute", {}, cache_key="test_key") + assert res1.success is True + assert res1.data == "Expensive Compute Result" + assert call_count == 1 + + # Second call uses cache + res2 = await registry.execute_with_resilience("compute", {}, cache_key="test_key") + assert res2.success is True + assert res2.data == "Expensive Compute Result" + assert call_count == 1 # count did not increase + assert "cache_hit" in res2.stages_executed