tests_python/test_model_catalog.py

import json

from arena.model_catalog import ModelCatalog, ModelProviderConfig, eligible_models


class FakeCatalog(ModelCatalog):
    def __init__(self, payload):
        super().__init__(ttl_seconds=60)
        self.payload = payload
        self.urls: list[str] = []

    def _get_json(self, url: str, *, api_key: str):
        self.urls.append(url)
        return self.payload


def test_local_catalog_without_endpoint_returns_cached_zerogpu_model() -> None:
    catalog = ModelCatalog()

    models = catalog.list_models(ModelProviderConfig(provider="local", model="local-model"))

    assert [model.id for model in models] == ["local-model"]
    assert models[0].provider == "local"


def test_local_catalog_cache_is_keyed_by_model_id() -> None:
    catalog = ModelCatalog()

    first = catalog.list_models(ModelProviderConfig(provider="local", model="local-7b"))
    second = catalog.list_models(ModelProviderConfig(provider="local", model="local-12b"))

    assert [model.id for model in first] == ["local-7b"]
    assert [model.id for model in second] == ["local-12b"]


def test_openai_compatible_catalog_normalizes_models_and_caches_without_user_key() -> None:
    catalog = FakeCatalog(
        {
            "data": [
                {
                    "id": "qwen/qwen2.5-7b-instruct",
                    "name": "Qwen 2.5 7B Instruct",
                    "providers": [
                        {
                            "provider": "deepinfra",
                            "status": "live",
                            "context_length": 32768,
                            "supports_tools": True,
                            "pricing": {"input": "0.1", "output": "0.2"},
                        }
                    ],
                },
                {
                    "id": "meta-llama/llama-3.1-70b-instruct",
                    "name": "Llama 3.1 70B Instruct",
                },
            ]
        }
    )
    config = ModelProviderConfig(provider="openrouter")

    first = catalog.list_models(config)
    second = catalog.list_models(config)

    assert first == second
    assert len(catalog.urls) == 1
    assert first[0].id == "qwen/qwen2.5-7b-instruct"
    assert first[0].label == "Qwen 2.5 7B Instruct"
    assert first[0].context_length == 32768
    assert first[0].supports_tools is True
    assert first[0].input_price == 0.1
    assert first[0].parameter_count_b == 7.0
    assert first[0].eligible is True
    assert first[1].parameter_count_b == 70.0
    assert first[1].eligible is False
    assert [model.id for model in eligible_models(first)] == ["qwen/qwen2.5-7b-instruct"]


def test_openai_compatible_catalog_detects_parameter_count_from_provider_metadata() -> None:
    catalog = FakeCatalog(
        {
            "data": [
                {
                    "id": "nvidia/nemotron-3.5-content-safety:free",
                    "name": "NVIDIA: Nemotron 3.5 Content Safety (free)",
                    "hugging_face_id": "nvidia/Nemotron-3.5-Content-Safety",
                    "description": "A compact 4B-parameter multimodal guardrail model.",
                },
                {
                    "id": "nvidia/nemotron-3-ultra",
                    "name": "NVIDIA: Nemotron 3 Ultra",
                    "description": "55B active parameters out of 550B total.",
                },
            ]
        }
    )

    models = catalog.list_models(ModelProviderConfig(provider="openrouter"))

    assert models[0].parameter_count_b == 4.0
    assert models[0].eligible is True
    assert models[1].parameter_count_b == 605.0
    assert models[1].eligible is False
    assert [model.id for model in eligible_models(models)] == ["nvidia/nemotron-3.5-content-safety:free"]


def test_user_key_catalog_fetch_is_cached_until_refresh() -> None:
    catalog = FakeCatalog({"data": [{"id": "qwen/qwen2.5-7b-instruct"}]})
    config = ModelProviderConfig(provider="openrouter", api_key="sk-user")

    catalog.list_models(config)
    catalog.list_models(config)

    assert len(catalog.urls) == 1

    catalog.refresh_models(config)

    assert len(catalog.urls) == 2


def test_gemini_catalog_filters_generate_content_models() -> None:
    catalog = FakeCatalog(
        {
            "models": [
                {
                    "name": "models/gemini-2.5-flash",
                    "supportedGenerationMethods": ["generateContent"],
                    "inputTokenLimit": 1048576,
                },
                {
                    "name": "models/embedding-001",
                    "supportedGenerationMethods": ["embedContent"],
                },
            ]
        }
    )

    models = catalog.list_models(ModelProviderConfig(provider="gemini", api_key="key"))

    assert [model.id for model in models] == ["gemini-2.5-flash"]
    assert models[0].context_length == 1048576
    assert models[0].eligible is False
    assert models[0].status == "unknown_parameters"


def test_model_provider_config_public_selection_drops_secret() -> None:
    config = ModelProviderConfig(
        provider="openrouter",
        model="qwen/model",
        api_key="sk-secret",
        base_url="https://example.test/v1",
    )

    public_json = config.public_selection().model_dump_json()

    assert "sk-secret" not in public_json
    assert json.loads(public_json) == {
        "provider": "openrouter",
        "model": "qwen/model",
        "base_url": "https://example.test/v1",
    }