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", }