from __future__ import annotations import time from dataclasses import dataclass from datetime import date, timedelta from pathlib import Path import httpx import tomlkit import yaml @dataclass class ProbeResult: provider: str model: str key_index: int status: str http_code: int | None message: str response_time_ms: float refresh_url: str | None = None def classify_response( provider: str, model: str, key_index: int, response: httpx.Response | None = None, *, error: Exception | None = None, ) -> ProbeResult: if error is not None: return ProbeResult( provider=provider, model=model, key_index=key_index, status="dead", http_code=None, message=str(error), response_time_ms=0, ) if response is None: raise ValueError("classify_response requires either a response or an error") code = response.status_code try: # nosemgrep: python-no-try-except-anywhere -- response.json() can fail on non-JSON bodies from diverse providers raw_body = response.json() except Exception: raw_body = {} body = raw_body if isinstance(raw_body, dict) else {} if code == 200: if str(body.get("status", "")) == "439": return ProbeResult( provider=provider, model=model, key_index=key_index, status="expired", http_code=code, message=body.get("msg", "expired"), response_time_ms=0, ) return ProbeResult( provider=provider, model=model, key_index=key_index, status="alive", http_code=code, message="OK", response_time_ms=0, ) if code in (401, 403): return ProbeResult( provider=provider, model=model, key_index=key_index, status="expired", http_code=code, message=_extract_error_message(body), response_time_ms=0, ) if code == 402: return ProbeResult( provider=provider, model=model, key_index=key_index, status="exhausted", http_code=code, message=_extract_error_message(body), response_time_ms=0, ) if code == 429: msg = _extract_error_message(body) status = "exhausted" if "quota" in msg.lower() else "rate_limited" return ProbeResult( provider=provider, model=model, key_index=key_index, status=status, http_code=code, message=msg, response_time_ms=0, ) return ProbeResult( provider=provider, model=model, key_index=key_index, status="dead", http_code=code, message=_extract_error_message(body), response_time_ms=0, ) def _extract_error_message(body: dict | list | str) -> str: if not isinstance(body, dict): return str(body)[:200] error = body.get("error", "") if isinstance(error, dict): return error.get("message", str(error)) return str(error) def hot_remove_deployment( proxy_url: str, model_id: str, *, master_key: str, ) -> bool: """Remove a deployment from the running LiteLLM proxy. Requires LiteLLM to be connected to a database (Postgres/SQLite). Without a DB, this call is silently ignored by LiteLLM. Returns True if the API call succeeded (200), False otherwise. """ with httpx.Client() as client: response = client.post( f"{proxy_url}/model/delete", json={"id": model_id}, headers={"Authorization": f"Bearer {master_key}"}, ) return response.status_code == 200 def probe_deployment( base_url: str, model_id: str, api_key: str, provider: str, key_index: int, *, timeout: float = 30.0, ) -> ProbeResult: try: with httpx.Client(timeout=timeout) as client: start = time.monotonic() response = client.post( f"{base_url}/chat/completions", json={ "model": model_id, "messages": [{"role": "user", "content": "Say hi"}], "max_tokens": 5, }, headers={"Authorization": f"Bearer {api_key}"}, ) elapsed_ms = (time.monotonic() - start) * 1000 result = classify_response(provider, model_id, key_index, response) result.response_time_ms = elapsed_ms return result except (httpx.ConnectTimeout, httpx.ReadTimeout, httpx.ConnectError) as exc: return classify_response(provider, model_id, key_index, error=exc) def update_secrets_on_probe_results( results: list[ProbeResult], secrets_path: Path, providers_path: Path, *, proxy_url: str | None = None, master_key: str | None = None, ) -> None: secrets_doc = tomlkit.loads(secrets_path.read_text()) providers_data = yaml.safe_load(providers_path.read_text()) providers = providers_data.get("providers", {}) modified = False alive_providers: set[str] = {r.provider for r in results if r.status in ("alive", "rate_limited")} for result in results: if result.status in ("alive", "rate_limited"): continue provider_name = result.provider # "dead" means the model is unavailable, not necessarily the key. # Only disable the provider if zero models are alive for it. if result.status == "dead" and provider_name in alive_providers: if proxy_url is not None: model_id = f"{provider_name}::{result.model}::{result.key_index}" hot_remove_deployment(proxy_url, model_id, master_key=master_key or "") continue if result.status == "exhausted": provider_config = providers.get(provider_name, {}) if provider_config.get("reset_schedule") == "monthly": disabled_until = _next_first_of_month() else: disabled_until = date.today() + timedelta(days=1) elif result.status in ("expired", "dead"): disabled_until = date.today() + timedelta(days=1) else: continue if "keys" in secrets_doc and provider_name in secrets_doc["keys"]: secrets_doc["keys"][provider_name]["disabled_until"] = disabled_until modified = True if proxy_url is not None: model_id = f"{provider_name}::{result.model}::{result.key_index}" hot_remove_deployment(proxy_url, model_id, master_key=master_key or "") if modified: secrets_path.write_text(tomlkit.dumps(secrets_doc)) def _next_first_of_month() -> date: today = date.today() if today.month == 12: return date(today.year + 1, 1, 1) return date(today.year, today.month + 1, 1)