model-selector/model_selector/providers.py

from __future__ import annotations

import asyncio
from dataclasses import dataclass
from typing import Any

import httpx

from .config import provider_label


_VALID_INPUT_MODALITIES = {"text", "image", "video", "audio"}


@dataclass(frozen=True)
class ProviderModel:
    provider_key: str
    provider_label: str
    model_id: str
    name: str
    source: str
    reasoning: bool | None
    context_window: int | None
    max_tokens: int | None
    pricing_prompt: str | None
    pricing_completion: str | None
    pricing_request: str | None
    input_modalities: tuple[str, ...]
    output_modalities: tuple[str, ...]
    supported_parameters: tuple[str, ...]
    active: bool | None
    features: tuple[str, ...]
    raw: dict[str, Any]
    default_parameters: dict[str, Any] | None = None

    @property
    def key(self) -> str:
        return f"{self.provider_key}/{self.model_id}"


def _feature_set(*values: str) -> tuple[str, ...]:
    seen: list[str] = []
    for value in values:
        if value and value not in seen:
            seen.append(value)
    return tuple(seen)


def _format_price(value: Any) -> str | None:
    if value is None:
        return None
    return str(value)


def infer_reasoning(raw: dict[str, Any], features: tuple[str, ...]) -> bool | None:
    if "reasoning" in raw:
        return bool(raw.get("reasoning"))
    if "reasoning" in features:
        return True
    default_parameters = raw.get("default_parameters")
    if isinstance(default_parameters, dict) and default_parameters.get("reasoning") is not None:
        return True
    return None


def provider_model_to_config_entry(model: ProviderModel) -> dict[str, Any]:
    input_modalities = _sanitize_input_modalities(model.input_modalities)
    entry = {
        "id": model.model_id,
        "name": model.name,
        "input": list(input_modalities),
        "contextWindow": model.context_window,
        "maxTokens": model.max_tokens,
    }
    if model.reasoning is True:
        entry["reasoning"] = True
    thinking = _openrouter_thinking_parameters(model)
    if thinking:
        entry["thinking"] = thinking
    return entry


def _sanitize_input_modalities(modalities: tuple[str, ...] | list[str] | None) -> tuple[str, ...]:
    filtered = tuple(modality for modality in (modalities or ()) if modality in _VALID_INPUT_MODALITIES)
    return filtered or ("text",)


def _openrouter_thinking_parameters(model: ProviderModel) -> dict[str, Any] | None:
    if model.provider_key != "openrouter":
        return None
    default_parameters = model.default_parameters or {}
    for key in ("reasoning", "thinking"):
        value = default_parameters.get(key)
        if isinstance(value, dict) and value:
            return dict(value)
    return None


def derive_features(provider_key: str, raw: dict[str, Any]) -> tuple[str, ...]:
    features: list[str] = []

    pricing = raw.get("pricing") or {}
    zero_pricing = False
    if pricing:
        price_values = [pricing.get("prompt"), pricing.get("completion"), pricing.get("request")]
        zero_pricing = all(str(item) in {"0", "0.0", "0.00000", "0.000000"} for item in price_values if item is not None)
    if raw.get("id", "").endswith(":free") or zero_pricing:
        features.append("free")

    supported_parameters = raw.get("supported_parameters") or []
    if any(param in supported_parameters for param in ("response_format", "json_schema")):
        features.append("structured output")
    if "tools" in supported_parameters:
        features.append("tool use")

    architecture = raw.get("architecture") or {}
    if architecture.get("input_modalities"):
        for modality in architecture.get("input_modalities", []):
            features.append(str(modality))
    if architecture.get("output_modalities"):
        for modality in architecture.get("output_modalities", []):
            features.append(str(modality))

    provider_hints = {
        "groq-cloud": {
            "openai/gpt-oss-20b": [
                "structured output",
                "browser search",
                "code execution",
                "reasoning",
            ],
            "openai/gpt-oss-120b": [
                "structured output",
                "browser search",
                "code execution",
                "reasoning",
            ],
            "openai/gpt-oss-safeguard-20b": [
                "structured output",
                "safety",
            ],
            "meta-llama/llama-4-scout-17b-16e-instruct": [
                "structured output",
                "reasoning",
            ],
        },
        "openai": {
            "gpt-5.4": ["structured output", "reasoning"],
            "gpt-5.4-mini": ["structured output", "reasoning"],
            "gpt-5.4-nano": ["structured output", "reasoning"],
            "gpt-5.1-codex": ["structured output", "reasoning", "code"],
            "gpt-5.1-codex-mini": ["structured output", "reasoning", "code"],
            "gpt-5.3-chat-latest": ["structured output", "reasoning"],
            "gpt-5-nano": ["structured output", "reasoning"],
            "gpt-4.1-nano-2025-04-14": ["structured output"],
        },
    }
    features.extend(provider_hints.get(provider_key, {}).get(raw.get("id", ""), []))

    if provider_key == "openrouter":
        for param in supported_parameters:
            if param == "response_format":
                features.append("structured output")
            elif param in {"tools", "tool_choice"}:
                features.append("tool use")
            elif param in {"reasoning", "thinking"} or param.startswith("reasoning"):
                features.append("reasoning")
            elif param == "seed":
                features.append("seed")

    if raw.get("active") is False:
        features.append("inactive")

    return _feature_set(*features)


class ProviderAdapter:
    provider_key: str
    base_url: str
    api_key_env: str
    model_path: str

    def __init__(self, base_url: str, api_key_env: str) -> None:
        self.base_url = base_url.rstrip("/")
        self.api_key_env = api_key_env

    def api_key(self) -> str:
        import os

        return os.environ.get(self.api_key_env, "").strip()

    def headers(self) -> dict[str, str]:
        headers = {"Content-Type": "application/json"}
        api_key = self.api_key()
        if api_key:
            headers["Authorization"] = f"Bearer {api_key}"
        return headers

    async def fetch(self) -> list[ProviderModel]:
        raise NotImplementedError

    async def _get_json(self, path: str) -> dict[str, Any]:
        async with httpx.AsyncClient(timeout=20.0) as client:
            response = await client.get(f"{self.base_url}{path}", headers=self.headers())
            response.raise_for_status()
            return response.json()


class OpenRouterAdapter(ProviderAdapter):
    provider_key = "openrouter"
    model_path = "/models"

    async def fetch(self) -> list[ProviderModel]:
        if not self.api_key():
            raise RuntimeError("missing OPENROUTER_API_KEY")
        payload = await self._get_json(self.model_path)
        models: list[ProviderModel] = []
        for item in payload.get("data", []):
            pricing = item.get("pricing") or {}
            architecture = item.get("architecture") or {}
            default_parameters = item.get("default_parameters") or {}
            raw = dict(item)
            raw["supported_parameters"] = item.get("supported_parameters") or []
            raw["pricing"] = pricing
            raw["architecture"] = architecture
            raw["default_parameters"] = default_parameters
            features = derive_features(self.provider_key, raw)
            models.append(
                ProviderModel(
                    provider_key=self.provider_key,
                    provider_label=provider_label(self.provider_key),
                    model_id=str(item.get("id")),
                    name=str(item.get("name", item.get("id"))),
                    source="live",
                    reasoning=infer_reasoning(raw, features),
                    context_window=item.get("context_length"),
                    max_tokens=(item.get("top_provider") or {}).get("max_completion_tokens"),
                    pricing_prompt=_format_price(pricing.get("prompt")),
                    pricing_completion=_format_price(pricing.get("completion")),
                    pricing_request=_format_price(pricing.get("request")),
                    input_modalities=_sanitize_input_modalities(architecture.get("input_modalities")),
                    output_modalities=tuple(architecture.get("output_modalities") or ()),
                    supported_parameters=tuple(item.get("supported_parameters") or ()),
                    active=None,
                    features=features,
                    raw=raw,
                    default_parameters=default_parameters if isinstance(default_parameters, dict) else None,
                )
            )
        return models


class GroqAdapter(ProviderAdapter):
    provider_key = "groq-cloud"
    model_path = "/models"

    async def fetch(self) -> list[ProviderModel]:
        if not self.api_key():
            raise RuntimeError("missing GROQ_API_KEY")
        payload = await self._get_json(self.model_path)
        models: list[ProviderModel] = []
        for item in payload.get("data", []):
            raw = dict(item)
            features = derive_features(self.provider_key, raw)
            models.append(
                ProviderModel(
                    provider_key=self.provider_key,
                    provider_label=provider_label(self.provider_key),
                    model_id=str(item.get("id")),
                    name=str(item.get("id")),
                    source="live",
                    reasoning=infer_reasoning(raw, features),
                    context_window=item.get("context_window"),
                    max_tokens=None,
                    pricing_prompt=None,
                    pricing_completion=None,
                    pricing_request=None,
                    input_modalities=("text",),
                    output_modalities=("text",),
                    supported_parameters=(),
                    active=item.get("active"),
                    features=features,
                    raw=raw,
                )
            )
        return models


class OpenAIAdapter(ProviderAdapter):
    provider_key = "openai"
    model_path = "/models"

    async def fetch(self) -> list[ProviderModel]:
        if not self.api_key():
            raise RuntimeError("missing OPENAI_API_KEY")
        payload = await self._get_json(self.model_path)
        models: list[ProviderModel] = []
        for item in payload.get("data", []):
            raw = dict(item)
            features = derive_features(self.provider_key, raw)
            models.append(
                ProviderModel(
                    provider_key=self.provider_key,
                    provider_label=provider_label(self.provider_key),
                    model_id=str(item.get("id")),
                    name=str(item.get("id")),
                    source="live",
                    reasoning=infer_reasoning(raw, features),
                    context_window=None,
                    max_tokens=None,
                    pricing_prompt=None,
                    pricing_completion=None,
                    pricing_request=None,
                    input_modalities=("text",),
                    output_modalities=("text",),
                    supported_parameters=(),
                    active=None,
                    features=features,
                    raw=raw,
                )
            )
        return models


def build_provider_adapters() -> dict[str, ProviderAdapter]:
    return {
        "openai": OpenAIAdapter("https://api.openai.com/v1", "OPENAI_API_KEY"),
        "groq-cloud": GroqAdapter("https://api.groq.com/openai/v1", "GROQ_API_KEY"),
        "openrouter": OpenRouterAdapter("https://openrouter.ai/api/v1", "OPENROUTER_API_KEY"),
    }


async def fetch_with_timeout(adapter: ProviderAdapter) -> list[ProviderModel]:
    return await asyncio.wait_for(adapter.fetch(), timeout=30.0)