benchmarks/model_watch_adapter.py

"""Observe-only model catalog adapters.

This module performs discovery only. It never updates ai_providers, selects a
fallback, or persists credentials. Callers can use CatalogResult as an audit
record and decide separately whether a later approval/apply phase is allowed.
"""
from __future__ import annotations

from dataclasses import dataclass, field, replace
import asyncio
from enum import Enum
import json
import os
import re
from typing import Any, Awaitable, Callable, Mapping, Optional

import httpx


class CatalogStatus(str, Enum):
    AVAILABLE = "available"
    UNAUTHORIZED = "unauthorized"
    FORBIDDEN = "forbidden"
    RATE_LIMITED = "rate_limited"
    PROVIDER_ERROR = "provider_error"
    TIMEOUT = "timeout"
    NETWORK_ERROR = "network_error"
    MALFORMED = "malformed"


@dataclass(frozen=True)
class ModelWatchConfig:
    """Safety gate for optional model updates.

    Discovery remains observe-only by default. Auto-apply is enabled only when
    the explicit flag and approval marker are both present; callers must also
    provide an allowlist of provider/old/new model triples.
    """

    auto_apply_enabled: bool = False
    approval_marker: str = ""
    required_approval_marker: str = "I_UNDERSTAND_MODEL_UPDATES"
    approved_updates: tuple[tuple[str, str, str], ...] = ()

    @classmethod
    def from_env(cls) -> "ModelWatchConfig":
        raw_updates = os.getenv("MODEL_AUTO_APPLY_ALLOWLIST", "")
        updates: list[tuple[str, str, str]] = []
        for item in raw_updates.split(","):
            parts = tuple(part.strip() for part in item.split("|"))
            if len(parts) == 3 and all(parts):
                updates.append(parts)  # type: ignore[arg-type]
        return cls(
            auto_apply_enabled=os.getenv("MODEL_AUTO_APPLY_ENABLED", "0").lower() in {"1", "true", "yes"},
            approval_marker=os.getenv("MODEL_AUTO_APPLY_APPROVAL", ""),
            approved_updates=tuple(updates),
        )

    @property
    def can_auto_apply(self) -> bool:
        return (
            self.auto_apply_enabled
            and self.approval_marker == self.required_approval_marker
            and bool(self.approved_updates)
        )


@dataclass(frozen=True)
class ProviderProfile:
    provider: str
    profile: str
    base_url: str
    api_key: str
    default_model: str
    auth_mode: str = "bearer"  # bearer | query_key


@dataclass(frozen=True)
class CatalogResult:
    provider: str
    profile: str
    status: CatalogStatus
    http_status: Optional[int] = None
    models: tuple[str, ...] = ()
    default_available: Optional[bool] = None
    retry_after_seconds: Optional[int] = None
    detail: str = ""
    checked_url: str = ""
    metadata: Mapping[str, Any] = field(default_factory=dict)

    @property
    def should_auto_apply(self) -> bool:
        """Observe-only invariant: this adapter can never authorize a write."""
        return False

    def as_audit_record(self) -> dict[str, Any]:
        return {
            "provider": self.provider,
            "profile": self.profile,
            "status": self.status.value,
            "http_status": self.http_status,
            "model_count": len(self.models),
            "default_available": self.default_available,
            "retry_after_seconds": self.retry_after_seconds,
            "detail": self.detail[:240],
            "checked_url": self.checked_url,
            "metadata": dict(self.metadata),
        }


_RETRY_AFTER_SECONDS = re.compile(r"^\s*(\d+)\s*$")


def models_url(base_url: str) -> str:
    """Normalize common OpenAI-compatible base URLs to a models endpoint."""
    value = base_url.rstrip("/")
    for suffix in ("/chat/completions", "/completions"):
        if value.endswith(suffix):
            value = value[: -len(suffix)]
    if not value.endswith("/models"):
        value += "/models"
    return value


def _retry_after(headers: Mapping[str, str]) -> Optional[int]:
    raw = headers.get("retry-after") or headers.get("Retry-After")
    if not raw:
        return None
    match = _RETRY_AFTER_SECONDS.match(raw)
    return int(match.group(1)) if match else None


def _safe_detail(response: httpx.Response) -> str:
    """Return bounded provider detail without authorization headers or secrets."""
    try:
        payload = response.json()
        if isinstance(payload, Mapping):
            for key in ("error", "message", "detail", "code"):
                value = payload.get(key)
                if value is not None:
                    return str(value)[:240]
        return json.dumps(payload, ensure_ascii=True)[:240]
    except Exception:
        return response.text[:240]


def _models_from_gemini_payload(payload: Any) -> tuple[str, ...] | None:
    """Parse Gemini's native {models: [{name: 'models/<id>'}]} payload."""
    if not isinstance(payload, Mapping) or not isinstance(payload.get("models"), list):
        return None
    models: list[str] = []
    for item in payload["models"]:
        if not isinstance(item, Mapping):
            continue
        name = item.get("name") or item.get("baseModelId")
        if isinstance(name, str) and name.strip():
            normalized = name.strip()
            if normalized.startswith("models/"):
                normalized = normalized[len("models/"):]
            models.append(normalized)
    return tuple(dict.fromkeys(models))


def _models_from_payload(payload: Any) -> tuple[str, ...] | None:
    if isinstance(payload, Mapping):
        items = payload.get("data")
    else:
        items = payload
    if not isinstance(items, list):
        return None
    models: list[str] = []
    for item in items:
        if isinstance(item, Mapping) and isinstance(item.get("id"), str) and item["id"].strip():
            models.append(item["id"].strip())
    return tuple(dict.fromkeys(models))


@dataclass(frozen=True)
class ProfileScan:
    """Results plus profiles skipped because their provider returned 429."""
    results: tuple[CatalogResult, ...]
    skipped_rate_limited: tuple[CatalogResult, ...] = ()


class ObserveOnlyModelsAdapter:
    """Fetch a provider catalog and classify the result; never mutates state."""

    def __init__(
        self,
        *,
        timeout_seconds: float = 8.0,
        client: httpx.AsyncClient | None = None,
        config: ModelWatchConfig | None = None,
    ):
        self.timeout_seconds = timeout_seconds
        self._client = client
        self.config = config or ModelWatchConfig.from_env()

    @property
    def can_auto_apply(self) -> bool:
        """True only when every explicit safety gate is satisfied."""
        return self.config.can_auto_apply

    async def apply_updates(
        self,
        updates: list[tuple[str, str, str]],
        apply_callback: Callable[[str, str, str], Awaitable[None]],
    ) -> dict[str, Any]:
        """Apply only allowlisted updates through a caller-owned callback.

        The adapter never receives a database client and cannot mutate state on
        its own. With the default config this returns a dry-run result.
        """
        if not self.can_auto_apply:
            return {"applied": False, "dry_run": True, "reason": "auto_apply_disabled"}
        approved = set(self.config.approved_updates)
        applied = 0
        skipped = 0
        for provider, old_model, new_model in updates:
            if (provider, old_model, new_model) not in approved:
                skipped += 1
                continue
            await apply_callback(provider, old_model, new_model)
            applied += 1
        return {"applied": applied > 0, "dry_run": False, "applied_count": applied, "skipped_count": skipped}

    async def list_models(self, profile: ProviderProfile) -> CatalogResult:
        url = models_url(profile.base_url)
        headers = {"Accept": "application/json"}
        params: dict[str, str] = {}
        if profile.auth_mode == "query_key":
            params["key"] = profile.api_key
        elif profile.auth_mode != "none":
            headers["Authorization"] = f"Bearer {profile.api_key}"

        owns_client = self._client is None
        client = self._client or httpx.AsyncClient(timeout=self.timeout_seconds)
        try:
            response = await client.get(url, headers=headers, params=params)
            status = response.status_code
            if status == 401:
                return self._result(profile, url, CatalogStatus.UNAUTHORIZED, response)
            if status == 403:
                return self._result(profile, url, CatalogStatus.FORBIDDEN, response)
            if status == 429:
                return self._result(profile, url, CatalogStatus.RATE_LIMITED, response)
            if 500 <= status <= 599:
                return self._result(profile, url, CatalogStatus.PROVIDER_ERROR, response)
            if status != 200:
                return self._result(profile, url, CatalogStatus.NETWORK_ERROR, response)
            try:
                payload = response.json()
            except (ValueError, json.JSONDecodeError):
                return self._result(profile, url, CatalogStatus.MALFORMED, response, detail="invalid JSON")
            models = _models_from_payload(payload)
            if models is None:
                return self._result(profile, url, CatalogStatus.MALFORMED, response, detail="missing data list")
            return CatalogResult(
                provider=profile.provider,
                profile=profile.profile,
                status=CatalogStatus.AVAILABLE,
                http_status=status,
                models=models,
                default_available=profile.default_model in models,
                checked_url=url,
                detail="catalog fetched",
            )
        except httpx.TimeoutException as exc:
            return CatalogResult(profile.provider, profile.profile, CatalogStatus.TIMEOUT, detail=str(exc)[:240], checked_url=url)
        except httpx.RequestError as exc:
            return CatalogResult(profile.provider, profile.profile, CatalogStatus.NETWORK_ERROR, detail=str(exc)[:240], checked_url=url)
        finally:
            if owns_client:
                await client.aclose()

    @staticmethod
    def _result(profile: ProviderProfile, url: str, status: CatalogStatus, response: httpx.Response, *, detail: str = "") -> CatalogResult:
        return CatalogResult(
            provider=profile.provider,
            profile=profile.profile,
            status=status,
            http_status=response.status_code,
            retry_after_seconds=_retry_after(response.headers) if status == CatalogStatus.RATE_LIMITED else None,
            detail=detail or _safe_detail(response),
            checked_url=url,
        )


class GeminiModelsAdapter(ObserveOnlyModelsAdapter):
    """Observe-only adapter for Gemini's native ``models`` catalog."""

    async def list_models(self, profile: ProviderProfile) -> CatalogResult:
        url = models_url(profile.base_url)
        headers = {"Accept": "application/json"}
        params = {"key": profile.api_key} if profile.auth_mode != "none" else {}
        owns_client = self._client is None
        client = self._client or httpx.AsyncClient(timeout=self.timeout_seconds)
        try:
            response = await client.get(url, headers=headers, params=params)
            status = response.status_code
            if status == 401:
                return self._result(profile, url, CatalogStatus.UNAUTHORIZED, response)
            if status == 403:
                return self._result(profile, url, CatalogStatus.FORBIDDEN, response)
            if status == 429:
                return self._result(profile, url, CatalogStatus.RATE_LIMITED, response)
            if 500 <= status <= 599:
                return self._result(profile, url, CatalogStatus.PROVIDER_ERROR, response)
            if status != 200:
                return self._result(profile, url, CatalogStatus.NETWORK_ERROR, response)
            try:
                payload = response.json()
            except (ValueError, json.JSONDecodeError):
                return self._result(profile, url, CatalogStatus.MALFORMED, response, detail="invalid JSON")
            models = _models_from_gemini_payload(payload)
            if models is None:
                return self._result(profile, url, CatalogStatus.MALFORMED, response, detail="missing models list")
            return CatalogResult(
                provider=profile.provider,
                profile=profile.profile,
                status=CatalogStatus.AVAILABLE,
                http_status=status,
                models=models,
                default_available=profile.default_model in models,
                checked_url=url,
                detail="Gemini catalog fetched",
                metadata={"catalog_format": "gemini_native"},
            )
        except httpx.TimeoutException as exc:
            return CatalogResult(profile.provider, profile.profile, CatalogStatus.TIMEOUT, detail=str(exc)[:240], checked_url=url)
        except httpx.RequestError as exc:
            return CatalogResult(profile.provider, profile.profile, CatalogStatus.NETWORK_ERROR, detail=str(exc)[:240], checked_url=url)
        finally:
            if owns_client:
                await client.aclose()


async def scan_profiles(
    profiles: list[ProviderProfile],
    *,
    adapter: ObserveOnlyModelsAdapter | None = None,
) -> ProfileScan:
    """Scan a mixed pool and isolate 429 profiles without blocking healthy ones."""
    adapter = adapter or ObserveOnlyModelsAdapter()
    results = await asyncio.gather(*(adapter.list_models(profile) for profile in profiles))
    skipped = tuple(
        replace(result, metadata={"skipped": True, "skip_reason": "rate_limited"})
        for result in results
        if result.status == CatalogStatus.RATE_LIMITED
    )
    active = tuple(result for result in results if result.status != CatalogStatus.RATE_LIMITED)
    return ProfileScan(results=active, skipped_rate_limited=skipped)