gemini_limits.py

"""Shared rate-limit bookkeeping for free-tier LLM APIs (all providers).

Every agent step is one LLM request, and free tiers allow only a handful of
requests per minute (and per day). This module tracks usage per *endpoint*
(a "provider|model" key) across the whole process so that:

* we never exceed each endpoint's requests-per-minute budget,
* an endpoint that returns 429 is put on cooldown and another is used,
* an endpoint whose DAILY quota is gone, or whose key is rejected, is dropped,
* if everything is cooling down for too long we raise QuotaExhausted instead
  of sleeping for hours (the evaluation then resumes later from progress.json).

(The file keeps its old name, gemini_limits.py, so existing imports still work.)
"""
from __future__ import annotations

import os
import re
import threading
import time
from collections import defaultdict, deque

# Load keys from a local .env file if present (never overrides real environment
# variables, so Space secrets / shell exports always win). Needs python-dotenv.
try:
    from dotenv import load_dotenv

    load_dotenv()
except ImportError:
    pass

# Vertex Gemini models used for the agent brain AND for the multimodal tools.
# Unknown names are harmless: a 404 just drops that model for the session.
MODEL_CHAIN = [
    "gemini-3.8-flash",
    "gemini-3.7-flash",
    "gemini-3.6-flash",
    "gemini-3.5-flash",
    "gemini-3.5-flash-lite",
    "gemini-3.1-flash-lite",
]

# Default requests/minute for an endpoint that did not register its own.
RPM = int(os.getenv("VERTEX_GEMINI_RPM") or os.getenv("GEMINI_RPM") or "4")
# Longest we are willing to wait for any endpoint to free up.
MAX_WAIT = float(os.getenv("MAX_WAIT_SECONDS", "180"))


class QuotaExhausted(RuntimeError):
    """Every configured endpoint is out of quota or cooling down too long."""


_lock = threading.Lock()
_calls: dict[str, deque] = defaultdict(deque)
_blocked_until: dict[str, float] = defaultdict(float)
_dead: dict[str, str] = {}  # key -> reason
_rpm: dict[str, int] = {}     # budget per group
_group: dict[str, str] = {}   # key -> budget group (== key unless shared)


def make_key(provider: str, model: str) -> str:
    return f"{provider}|{model}"


def model_name(key: str) -> str:
    return key.split("|", 1)[1] if "|" in key else key


def register(key: str, rpm: int | None, group: str | None = None) -> None:
    """Set the requests/minute budget of `key`.

    `group` makes several keys share ONE budget and ONE 429 cooldown (e.g. all
    NVIDIA NIM models share a single 40 requests/min allowance).
    """
    grp = group or key
    _group[key] = grp
    if rpm:
        _rpm[grp] = int(rpm)


def gemini_models() -> list[str]:
    """Endpoint keys for Vertex Gemini (multimodal tools use these too)."""
    keys = [make_key("gemini", m) for m in model_list()]
    for k in keys:
        register(k, RPM)
    return keys


def model_list() -> list[str]:
    """Return the Vertex Gemini chain, optionally overridden by environment."""
    configured = os.getenv("VERTEX_GEMINI_MODELS", "").strip()
    if configured:
        return [model.strip() for model in configured.split(",") if model.strip()]
    return list(MODEL_CHAIN)


def parse_retry_delay(text: str, default: float = 20.0) -> float:
    """Understands '4.07s', '7m12.5s', '1h2m', retryDelay: '4s', retry-after."""
    m = re.search(r"(?:retry|try again) in\s+((?:\d+(?:\.\d+)?\s*[hms]\s*)+)", text, re.I)
    if not m:
        m2 = re.search(r"retryDelay['\"]?:\s*['\"]?(\d+(?:\.\d+)?)s", text)
        return float(m2.group(1)) if m2 else default
    total = 0.0
    for num, unit in re.findall(r"(\d+(?:\.\d+)?)\s*([hms])", m.group(1), re.I):
        total += float(num) * {"h": 3600, "m": 60, "s": 1}[unit.lower()]
    return total or default


def wait_time(key: str) -> float:
    """Seconds until `key` may be called again (0 = now)."""
    now = time.time()
    grp = _group.get(key, key)
    q = _calls[grp]
    while q and now - q[0] >= 60:
        q.popleft()
    wait = max(0.0, _blocked_until[key] - now)
    limit = _rpm.get(grp, RPM)
    if len(q) >= limit:
        wait = max(wait, 60 - (now - q[0]) + 0.5)
    return wait


def acquire(keys: list[str], max_wait: float | None = None) -> str:
    """Block until some endpoint in `keys` is usable; reserve a slot on it.

    Preference order = order of `keys` (first one with spare capacity wins).
    """
    limit = MAX_WAIT if max_wait is None else max_wait
    while True:
        with _lock:
            live = [k for k in keys if k not in _dead]
            if not live:
                reasons = "; ".join(f"{k}: {_dead.get(k, '?')}" for k in keys)
                raise QuotaExhausted(
                    "No usable model left (" + reasons + "). Wait for the quota "
                    "to reset, add another provider key, or enable billing."
                )
            waits = {k: wait_time(k) for k in live}
            best = min(live, key=lambda k: waits[k])  # ties -> earliest in list
            if waits[best] <= 0:
                _calls[_group.get(best, best)].append(time.time())
                return best
            delay = waits[best]
            if delay > limit:
                raise QuotaExhausted(
                    f"Every model is rate-limited; the earliest frees up in "
                    f"{delay / 60:.1f} min. Re-run later: solved questions are "
                    "saved and will be skipped."
                )
        print(f"[RATE] All models busy, waiting {delay:.0f}s...")
        time.sleep(min(delay, 30))


def _has_code(low: str, *codes: str) -> bool:
    """True if an HTTP status code appears as a whole number (not inside 14012)."""
    return any(re.search(rf"(?<![\d.]){c}(?![\d])", low) for c in codes)


def report_error(key: str, exc: BaseException) -> bool:
    """Record a failure. Returns True if the caller should try another endpoint."""
    text = str(exc)
    low = text.lower()
    squashed = low.replace("_", "").replace(" ", "")
    status = getattr(exc, "status_code", None)
    if status:  # openai exceptions carry the real status; make sure it's matched
        low += f" {status}"
    with _lock:
        # --- permanent for this session ------------------------------------
        if any(s in squashed for s in (
            "perday", "tokensperday", "requestsperday", "dailylimit", "dailyquota",
        )):
            print(f"[RATE] {key}: daily quota exhausted.")
            _dead[key] = "daily quota exhausted"
            return True
        if (_has_code(low, "401", "403") or any(s in low for s in (
            "invalid api key", "incorrect api key", "api key not valid",
            "permission denied", "unauthorized", "authentication",
        ))) and not _has_code(low, "429"):
            print(f"[RATE] {key}: key rejected, dropping it.")
            _dead[key] = "API key rejected"
            return True
        if (_has_code(low, "404") or "not found" in low or "not_found" in low
                or "does not exist" in low or "no such model" in low):
            print(f"[RATE] {key}: model not available on this key, skipping.")
            _dead[key] = "model not available"
            return True
        # --- temporary -------------------------------------------------------
        if _has_code(low, "429") or "resourceexhausted" in squashed or "ratelimit" in squashed \
                or "toomanyrequests" in squashed or "quota" in low:
            delay = parse_retry_delay(text) + 1
            grp = _group.get(key, key)
            for k, g in list(_group.items()) + [(key, grp)]:
                if g == grp:  # a shared budget cools down together
                    _blocked_until[k] = time.time() + delay
            print(f"[RATE] {key}: 429, cooling down {delay:.0f}s.")
            return True
        if any(s in low for s in (
            "tool_use_failed", "failed to call a function", "failed_generation",
            "request too large", "context length", "maximum context",
        )) or _has_code(low, "413"):
            _blocked_until[key] = time.time() + 2
            print(f"[RATE] {key}: request rejected ({text[:80]!r}), trying another.")
            return True
        if _has_code(low, "500", "502", "503", "504") or any(s in low for s in (
                "unavailable", "overloaded", "timeout", "timed out", "connection")):
            _blocked_until[key] = time.time() + 15
            return True
        if _has_code(low, "400", "422"):
            # Often a provider-specific quirk (tool_choice, schema...). Don't
            # kill the endpoint for good: park it briefly and try the others.
            _blocked_until[key] = time.time() + 30
            print(f"[RATE] {key}: bad request ({text[:100]!r}), trying another.")
            return True
    return False