scripts/probe_providers.py

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)