feat: gemini_guard — rate limiter + circuit breaker på alle Gemini-kall
This commit is contained in:
parent
6ffe23943c
commit
434f591f9c
|
|
@ -17,6 +17,7 @@ from typing import Literal
|
||||||
from google.adk.agents import Agent
|
from google.adk.agents import Agent
|
||||||
from google.adk.tools import FunctionTool
|
from google.adk.tools import FunctionTool
|
||||||
from google.adk.runners import Runner
|
from google.adk.runners import Runner
|
||||||
|
from gemini_guard import guarded_run, CircuitOpen, RateLimitExceeded
|
||||||
from google.adk.sessions.in_memory_session_service import InMemorySessionService
|
from google.adk.sessions.in_memory_session_service import InMemorySessionService
|
||||||
from google.genai import types
|
from google.genai import types
|
||||||
|
|
||||||
|
|
@ -282,7 +283,7 @@ async def _run_async(message: str, user_id: str, session_id: str, mode: str, cal
|
||||||
output_tokens = 0
|
output_tokens = 0
|
||||||
request_id = str(uuid.uuid4())
|
request_id = str(uuid.uuid4())
|
||||||
|
|
||||||
async for event in runner.run_async(user_id=user_id, session_id=session.id, new_message=new_message):
|
async for event in guarded_run(runner, user_id=user_id, session_id=session.id, new_message=new_message):
|
||||||
if event.is_final_response() and event.content and event.content.parts:
|
if event.is_final_response() and event.content and event.content.parts:
|
||||||
final_text = event.content.parts[0].text or ""
|
final_text = event.content.parts[0].text or ""
|
||||||
if hasattr(event, "usage_metadata") and event.usage_metadata:
|
if hasattr(event, "usage_metadata") and event.usage_metadata:
|
||||||
|
|
|
||||||
102
agents/core-logic/gemini_guard.py
Normal file
102
agents/core-logic/gemini_guard.py
Normal file
|
|
@ -0,0 +1,102 @@
|
||||||
|
"""
|
||||||
|
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
|
||||||
Loading…
Reference in New Issue
Block a user