diff --git a/agents/core-logic/agent.py b/agents/core-logic/agent.py index 331fc3d..76ccc4c 100644 --- a/agents/core-logic/agent.py +++ b/agents/core-logic/agent.py @@ -17,6 +17,7 @@ from typing import Literal from google.adk.agents import Agent from google.adk.tools import FunctionTool 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.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 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: final_text = event.content.parts[0].text or "" if hasattr(event, "usage_metadata") and event.usage_metadata: diff --git a/agents/core-logic/gemini_guard.py b/agents/core-logic/gemini_guard.py new file mode 100644 index 0000000..bbe6246 --- /dev/null +++ b/agents/core-logic/gemini_guard.py @@ -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