Chris Christiansen
9fcb9c354a
Some checks are pending
Check Python Version Consistency / Check Python Version (push) Waiting to run
103 lines
3.3 KiB
Python
103 lines
3.3 KiB
Python
"""
|
|
gemini_guard.py — Rate limiter + circuit breaker for alle Gemini-kall i OSVauco.
|
|
Wrapper rundt runner.run_async() i _run_async().
|
|
"""
|
|
import time
|
|
import asyncio
|
|
import logging
|
|
from collections import deque
|
|
from typing import AsyncIterator
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# --- Konfigurasjon (env-override mulig) ---
|
|
import os
|
|
RATE_LIMIT_RPM = int(os.environ.get("GUARD_RATE_LIMIT_RPM", "30"))
|
|
CIRCUIT_FAIL_THRESH = int(os.environ.get("GUARD_CIRCUIT_FAILS", "5"))
|
|
CIRCUIT_RESET_SECS = int(os.environ.get("GUARD_CIRCUIT_RESET_SECS", "60"))
|
|
|
|
class CircuitOpen(Exception):
|
|
pass
|
|
|
|
class RateLimitExceeded(Exception):
|
|
pass
|
|
|
|
class GeminiGuard:
|
|
"""
|
|
Én singleton per prosess.
|
|
- Rate limit: maks RATE_LIMIT_RPM kall per minutt (sliding window)
|
|
- Circuit breaker: åpner ved CIRCUIT_FAIL_THRESH feil på rad,
|
|
reset etter CIRCUIT_RESET_SECS sekunder
|
|
"""
|
|
def __init__(self):
|
|
self._timestamps: deque = deque()
|
|
self._failures = 0
|
|
self._opened_at = 0.0
|
|
self._state = "closed" # closed | open | half-open
|
|
|
|
def _check_rate(self):
|
|
now = time.monotonic()
|
|
cutoff = now - 60.0
|
|
while self._timestamps and self._timestamps[0] < cutoff:
|
|
self._timestamps.popleft()
|
|
if len(self._timestamps) >= RATE_LIMIT_RPM:
|
|
raise RateLimitExceeded(
|
|
f"Rate limit {RATE_LIMIT_RPM} RPM nådd ({len(self._timestamps)} kall siste minutt)"
|
|
)
|
|
self._timestamps.append(now)
|
|
|
|
def _check_circuit(self):
|
|
if self._state == "open":
|
|
elapsed = time.monotonic() - self._opened_at
|
|
if elapsed >= CIRCUIT_RESET_SECS:
|
|
self._state = "half-open"
|
|
logger.warning("[guard] circuit half-open — prøver igjen")
|
|
else:
|
|
raise CircuitOpen(
|
|
f"Circuit breaker åpen — venter {CIRCUIT_RESET_SECS - int(elapsed)}s"
|
|
)
|
|
|
|
def on_success(self):
|
|
self._failures = 0
|
|
self._state = "closed"
|
|
|
|
def on_failure(self, exc: Exception):
|
|
self._failures += 1
|
|
logger.error(f"[guard] Gemini-feil #{self._failures}: {exc}")
|
|
if self._failures >= CIRCUIT_FAIL_THRESH:
|
|
self._state = "open"
|
|
self._opened_at = time.monotonic()
|
|
logger.critical(f"[guard] Circuit ÅPNET etter {self._failures} feil")
|
|
|
|
async def run(self, runner, *, user_id, session_id, new_message) -> AsyncIterator:
|
|
self._check_circuit()
|
|
self._check_rate()
|
|
try:
|
|
events = []
|
|
async for event in runner.run_async(
|
|
user_id=user_id,
|
|
session_id=session_id,
|
|
new_message=new_message,
|
|
):
|
|
events.append(event)
|
|
self.on_success()
|
|
for e in events:
|
|
yield e
|
|
except (CircuitOpen, RateLimitExceeded):
|
|
raise
|
|
except Exception as exc:
|
|
self.on_failure(exc)
|
|
raise
|
|
|
|
_guard = GeminiGuard()
|
|
|
|
async def guarded_run(runner, *, user_id, session_id, new_message) -> AsyncIterator:
|
|
"""Drop-in erstatning for runner.run_async(...) """
|
|
async for event in _guard.run(
|
|
runner,
|
|
user_id=user_id,
|
|
session_id=session_id,
|
|
new_message=new_message,
|
|
):
|
|
yield event
|