"""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"(? 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