Spaces:
Sleeping
Sleeping
File size: 5,412 Bytes
1ebb69b | 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 | """
Nancy — Provider models.
Tracks per-provider runtime state: rate limits, circuit breaker status,
and routing metadata.
"""
from __future__ import annotations
import time
try:
from enum import StrEnum
except ImportError:
import enum
class StrEnum(str, enum.Enum):
pass
from typing import Any
from pydantic import BaseModel, Field
class CircuitState(StrEnum):
"""Circuit breaker states."""
CLOSED = "closed" # healthy — requests flow through
OPEN = "open" # tripped — requests are blocked
HALF_OPEN = "half_open" # cooldown expired — next request is a probe
class ProviderConfig(BaseModel):
"""
Static configuration for a single provider.
Loaded from the ``PROVIDERS_CONFIG`` environment variable or defaults.
"""
rpm: int = Field(default=10, description="Requests per minute limit.")
tpm: int = Field(default=40000, description="Tokens per minute limit.")
url_pattern: str = Field(
default="", description="Base URL pattern for the chatbot UI."
)
class ProviderState:
"""
Mutable runtime state for a single provider.
This is NOT a Pydantic model because it holds mutable counters and
timestamps that change on every request.
Attributes:
name: Provider identifier (e.g. ``"chatgpt"``).
config: Static provider configuration.
circuit_state: Current circuit breaker state.
consecutive_failures: Count of back-to-back failures.
last_failure_time: Timestamp of the most recent failure.
last_success_time: Timestamp of the most recent success.
request_timestamps: Rolling window of request timestamps for RPM.
"""
__slots__ = (
"name",
"config",
"circuit_state",
"consecutive_failures",
"last_failure_time",
"last_success_time",
"request_timestamps",
)
def __init__(self, name: str, config: ProviderConfig) -> None:
self.name = name
self.config = config
self.circuit_state = CircuitState.CLOSED
self.consecutive_failures: int = 0
self.last_failure_time: float = 0.0
self.last_success_time: float = 0.0
self.request_timestamps: list[float] = []
# ── Circuit Breaker ───────────────────────────────────────────────
def record_success(self) -> None:
"""Reset failure counter and close the circuit."""
self.consecutive_failures = 0
self.last_success_time = time.time()
self.circuit_state = CircuitState.CLOSED
def record_failure(self, threshold: int) -> None:
"""Increment failure counter; trip if threshold is reached."""
self.consecutive_failures += 1
self.last_failure_time = time.time()
if self.consecutive_failures >= threshold:
self.circuit_state = CircuitState.OPEN
def should_allow_request(self, threshold: int, cooldown: float) -> bool:
"""
Check whether the circuit breaker allows a request.
- CLOSED → always allow.
- OPEN → allow only if cooldown has elapsed (transition to HALF_OPEN).
- HALF_OPEN → allow (it's a probe request).
"""
if self.circuit_state == CircuitState.CLOSED:
return True
if self.circuit_state == CircuitState.OPEN:
elapsed = time.time() - self.last_failure_time
if elapsed >= cooldown:
self.circuit_state = CircuitState.HALF_OPEN
return True
return False
# HALF_OPEN — allow the probe
return True
# ── Rate Limiting (sliding window) ────────────────────────────────
def check_rate_limit(self) -> bool:
"""
Return True if the provider is within its RPM budget.
Prunes timestamps older than 60 seconds.
"""
now = time.time()
cutoff = now - 60.0
self.request_timestamps = [
ts for ts in self.request_timestamps if ts > cutoff
]
return len(self.request_timestamps) < self.config.rpm
def record_request(self) -> None:
"""Record a request timestamp for RPM tracking."""
self.request_timestamps.append(time.time())
# ── Serialization ─────────────────────────────────────────────────
def to_dict(self) -> dict[str, Any]:
"""Serialize for health / debug endpoints."""
now = time.time()
cutoff = now - 60.0
active_rpm = len([ts for ts in self.request_timestamps if ts > cutoff])
return {
"name": self.name,
"circuit_state": self.circuit_state.value,
"consecutive_failures": self.consecutive_failures,
"rpm_current": active_rpm,
"rpm_limit": self.config.rpm,
"last_failure_ago": (
round(now - self.last_failure_time, 1)
if self.last_failure_time
else None
),
"last_success_ago": (
round(now - self.last_success_time, 1)
if self.last_success_time
else None
),
}
|