generated from coulomb/repo-seed
66 lines
1.7 KiB
Python
66 lines
1.7 KiB
Python
import pytest
|
|
|
|
from llm_connect.rates import ModelRate, ModelRateRegistry
|
|
|
|
|
|
def test_default_registry_contains_openrouter_seed_models():
|
|
registry = ModelRateRegistry.default()
|
|
rates = registry.all()
|
|
|
|
assert len(rates) >= 9
|
|
assert rates["openai/gpt-4o-mini"].captured_at == "2026-05-17"
|
|
assert rates["openai/gpt-4o-mini"].source_url == "https://openrouter.ai/models"
|
|
|
|
|
|
def test_from_yaml_loads_package_shape(tmp_path):
|
|
path = tmp_path / "model-rates.yaml"
|
|
path.write_text(
|
|
"""
|
|
schema_version: 1
|
|
currency: USD
|
|
source_url: https://example.test/rates
|
|
captured_at: "2026-05-19"
|
|
rates:
|
|
vendor/model:
|
|
prompt_per_1k: 0.1
|
|
completion_per_1k: 0.2
|
|
""",
|
|
encoding="utf-8",
|
|
)
|
|
|
|
registry = ModelRateRegistry.from_yaml(path)
|
|
rate = registry.get("vendor/model")
|
|
|
|
assert rate == ModelRate(
|
|
model_id="vendor/model",
|
|
prompt_per_1k=0.1,
|
|
completion_per_1k=0.2,
|
|
currency="USD",
|
|
source_url="https://example.test/rates",
|
|
captured_at="2026-05-19",
|
|
)
|
|
|
|
|
|
def test_merged_with_overrides_matching_model():
|
|
base = ModelRateRegistry.default()
|
|
override = ModelRateRegistry(
|
|
{
|
|
"openai/gpt-4o-mini": ModelRate(
|
|
"openai/gpt-4o-mini",
|
|
prompt_per_1k=1,
|
|
completion_per_1k=2,
|
|
captured_at="override",
|
|
)
|
|
}
|
|
)
|
|
|
|
merged = base.merged_with(override)
|
|
|
|
assert merged.get("openai/gpt-4o-mini").prompt_per_1k == 1
|
|
assert merged.get("openai/gpt-4o-mini").captured_at == "override"
|
|
|
|
|
|
def test_negative_rates_are_rejected():
|
|
with pytest.raises(ValueError, match="prompt_per_1k"):
|
|
ModelRate("bad/model", prompt_per_1k=-1, completion_per_1k=0)
|