File size: 5,466 Bytes
31c8075 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 | 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",
}
|