Spaces:
Sleeping
Sleeping
File size: 5,698 Bytes
72bff80 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 | """
Circuit breaker pattern for external dependencies.
"""
import time
import threading
from enum import Enum
from typing import Callable, Any, Optional
from config import config
class CircuitState(Enum):
"""Circuit breaker states."""
CLOSED = "closed" # Normal operation
OPEN = "open" # Failing, reject requests
HALF_OPEN = "half_open" # Testing recovery
class CircuitBreakerError(Exception):
"""Raised when circuit is open."""
pass
class CircuitBreaker:
"""
Circuit breaker for external dependencies.
States:
- CLOSED: Normal operation, all requests pass through
- OPEN: Failure threshold reached, all requests rejected
- HALF_OPEN: After timeout, allow test requests
"""
def __init__(
self,
name: str,
failure_threshold: int = None,
timeout_seconds: int = None
):
"""
Initialize circuit breaker.
Args:
name: Name of the circuit (for logging)
failure_threshold: Number of consecutive failures before opening
timeout_seconds: Seconds to wait before attempting recovery
"""
self.name = name
self.failure_threshold = failure_threshold or config.CIRCUIT_BREAKER_FAILURE_THRESHOLD
self.timeout = timeout_seconds or config.CIRCUIT_BREAKER_TIMEOUT
self.state = CircuitState.CLOSED
self.failure_count = 0
self.last_failure_time = None
self.success_count = 0
self._lock = threading.Lock()
def call(self, func: Callable, *args, **kwargs) -> Any:
"""
Execute function with circuit breaker protection.
Args:
func: Function to call
*args, **kwargs: Arguments to pass to function
Returns:
Function result
Raises:
CircuitBreakerError: If circuit is open
"""
with self._lock:
# Check state transitions
if self.state == CircuitState.OPEN:
if self._should_attempt_reset():
self.state = CircuitState.HALF_OPEN
self.success_count = 0
else:
raise CircuitBreakerError(
f"Circuit breaker '{self.name}' is OPEN. "
f"Will retry after {self.timeout}s"
)
# Attempt the call
try:
result = func(*args, **kwargs)
self._on_success()
return result
except Exception as e:
self._on_failure()
raise e
def _should_attempt_reset(self) -> bool:
"""Check if enough time has passed to attempt reset."""
if self.last_failure_time is None:
return True
return time.time() - self.last_failure_time >= self.timeout
def _on_success(self):
"""Handle successful call."""
with self._lock:
self.failure_count = 0
if self.state == CircuitState.HALF_OPEN:
self.success_count += 1
# After 3 successful calls in HALF_OPEN, close the circuit
if self.success_count >= 3:
self.state = CircuitState.CLOSED
self.success_count = 0
def _on_failure(self):
"""Handle failed call."""
with self._lock:
self.failure_count += 1
self.last_failure_time = time.time()
# In HALF_OPEN, any failure immediately opens circuit
if self.state == CircuitState.HALF_OPEN:
self.state = CircuitState.OPEN
self.failure_count = 0
return
# In CLOSED, open after threshold
if self.failure_count >= self.failure_threshold:
self.state = CircuitState.OPEN
def reset(self):
"""Manually reset circuit breaker."""
with self._lock:
self.state = CircuitState.CLOSED
self.failure_count = 0
self.success_count = 0
self.last_failure_time = None
def get_state(self) -> dict:
"""Get current state for monitoring."""
with self._lock:
return {
"name": self.name,
"state": self.state.value,
"failure_count": self.failure_count,
"success_count": self.success_count,
"last_failure_time": self.last_failure_time
}
class CircuitBreakerManager:
"""Manages circuit breakers for different dependencies."""
def __init__(self):
self.breakers = {
"llm": CircuitBreaker("llm", failure_threshold=5, timeout_seconds=60),
"retriever": CircuitBreaker("retriever", failure_threshold=3, timeout_seconds=30),
"vector_store": CircuitBreaker("vector_store", failure_threshold=3, timeout_seconds=30),
}
def get_breaker(self, name: str) -> CircuitBreaker:
"""Get circuit breaker by name."""
if name not in self.breakers:
self.breakers[name] = CircuitBreaker(name)
return self.breakers[name]
def get_all_states(self) -> dict:
"""Get states of all circuit breakers."""
return {
name: breaker.get_state()
for name, breaker in self.breakers.items()
}
def reset_all(self):
"""Reset all circuit breakers."""
for breaker in self.breakers.values():
breaker.reset()
# Global circuit breaker manager
circuit_breaker_manager = CircuitBreakerManager()
|