mirror of
https://github.com/Routstr/routstr-core.git
synced 2026-07-30 15:26:14 +00:00
294 lines
9.4 KiB
Python
294 lines
9.4 KiB
Python
"""Tests for the model prioritization algorithm."""
|
|
|
|
import os
|
|
from types import SimpleNamespace
|
|
from unittest.mock import Mock
|
|
|
|
import pytest
|
|
|
|
# Set required env vars before importing
|
|
os.environ["UPSTREAM_BASE_URL"] = "http://test"
|
|
os.environ["UPSTREAM_API_KEY"] = "test"
|
|
|
|
from routstr.algorithm import ( # noqa: E402
|
|
calculate_model_cost_score,
|
|
create_model_mappings,
|
|
get_provider_penalty,
|
|
)
|
|
from routstr.payment.models import Architecture, Model, Pricing # noqa: E402
|
|
|
|
|
|
def create_test_model(
|
|
model_id: str,
|
|
prompt_price: float = 0.001,
|
|
completion_price: float = 0.002,
|
|
request_price: float = 0.0,
|
|
) -> Model:
|
|
"""Helper to create a test model with given pricing."""
|
|
return Model(
|
|
id=model_id,
|
|
name=f"Test {model_id}",
|
|
created=1234567890,
|
|
description="Test model",
|
|
context_length=8192,
|
|
architecture=Architecture(
|
|
modality="text",
|
|
input_modalities=["text"],
|
|
output_modalities=["text"],
|
|
tokenizer="gpt",
|
|
instruct_type=None,
|
|
),
|
|
pricing=Pricing(
|
|
prompt=prompt_price,
|
|
completion=completion_price,
|
|
request=request_price,
|
|
image=0.0,
|
|
web_search=0.0,
|
|
internal_reasoning=0.0,
|
|
),
|
|
)
|
|
|
|
|
|
def create_test_provider(
|
|
name: str,
|
|
base_url: str = "http://test.com",
|
|
*,
|
|
db_id: int | None = None,
|
|
models: list[Model] | None = None,
|
|
upstream_name: str | None = None,
|
|
) -> Mock:
|
|
"""Helper to create a test provider mock."""
|
|
provider = Mock()
|
|
provider.provider_type = name
|
|
provider.base_url = base_url
|
|
provider.db_id = db_id
|
|
provider.upstream_name = upstream_name or name
|
|
provider.get_cached_models.return_value = models or []
|
|
return provider
|
|
|
|
|
|
def test_calculate_model_cost_score_basic() -> None:
|
|
"""Test basic cost calculation."""
|
|
model = create_test_model("test-model", prompt_price=0.001, completion_price=0.002)
|
|
cost = calculate_model_cost_score(model)
|
|
|
|
# Expected: (1000 tokens * 0.001) + (500 tokens * 0.002) = 0.001 + 0.001 = 0.002
|
|
assert cost == 0.002
|
|
|
|
|
|
def test_calculate_model_cost_score_with_request_fee() -> None:
|
|
"""Test cost calculation with request fee."""
|
|
model = create_test_model(
|
|
"test-model",
|
|
prompt_price=0.001,
|
|
completion_price=0.002,
|
|
request_price=0.0005,
|
|
)
|
|
cost = calculate_model_cost_score(model)
|
|
|
|
# Expected: 0.001 + 0.001 + 0.0005 = 0.0025
|
|
assert cost == 0.0025
|
|
|
|
|
|
def test_calculate_model_cost_score_expensive_model() -> None:
|
|
"""Test cost calculation for expensive model."""
|
|
model = create_test_model(
|
|
"expensive-model", prompt_price=0.03, completion_price=0.06
|
|
)
|
|
cost = calculate_model_cost_score(model)
|
|
|
|
# Expected: (1000 * 0.03) + (500 * 0.06) = 0.03 + 0.03 = 0.06
|
|
assert cost == 0.06
|
|
|
|
|
|
def test_get_provider_penalty_regular_provider() -> None:
|
|
"""Test penalty for regular provider."""
|
|
provider = create_test_provider("regular-provider", "http://provider.com")
|
|
penalty = get_provider_penalty(provider)
|
|
assert penalty == 1.0
|
|
|
|
|
|
def test_get_provider_penalty_openrouter() -> None:
|
|
"""Test penalty for OpenRouter."""
|
|
provider = create_test_provider("openrouter", "https://openrouter.ai/api/v1")
|
|
penalty = get_provider_penalty(provider)
|
|
assert penalty == 1.001
|
|
|
|
|
|
def test_create_model_mappings_includes_db_override_for_missing_cached_model(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
"""Model overrides should still map when provider discovery misses the model."""
|
|
provider = create_test_provider(
|
|
"azure",
|
|
"https://example.openai.azure.com/openai/v1",
|
|
db_id=7,
|
|
models=[],
|
|
)
|
|
override_model = create_test_model("azure/gpt-4o")
|
|
override_model.canonical_slug = "azure-deployment"
|
|
|
|
def fake_row_to_model(*args, **kwargs) -> Model: # type: ignore[no-untyped-def]
|
|
return override_model
|
|
|
|
monkeypatch.setattr("routstr.payment.models._row_to_model", fake_row_to_model)
|
|
|
|
override_row = SimpleNamespace(id="azure/gpt-4o", upstream_provider_id=7, enabled=True)
|
|
|
|
model_instances, provider_map, unique_models = create_model_mappings(
|
|
upstreams=[provider],
|
|
overrides_by_key={("azure/gpt-4o", 7): (override_row, 1.01)},
|
|
disabled_model_keys=set(),
|
|
)
|
|
|
|
assert "azure/gpt-4o" in model_instances
|
|
assert [p for _, p in provider_map["azure/gpt-4o"]] == [provider]
|
|
assert "gpt-4o" in unique_models
|
|
|
|
|
|
def test_create_model_mappings_dedupes_with_provider_identity_not_provider_type(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
"""Different provider instances of same type should both survive dedupe."""
|
|
provider_a_model = create_test_model(
|
|
"azure/gpt-4o", prompt_price=0.01, completion_price=0.01
|
|
)
|
|
provider_a = create_test_provider(
|
|
"azure",
|
|
"https://a.openai.azure.com/openai/v1",
|
|
db_id=1,
|
|
models=[provider_a_model],
|
|
upstream_name="azure-a",
|
|
)
|
|
provider_b = create_test_provider(
|
|
"azure",
|
|
"https://b.openai.azure.com/openai/v1",
|
|
db_id=2,
|
|
models=[],
|
|
upstream_name="azure-b",
|
|
)
|
|
|
|
override_model = create_test_model(
|
|
"azure/gpt-4o", prompt_price=0.001, completion_price=0.001
|
|
)
|
|
override_model.canonical_slug = "azure-b-deployment"
|
|
|
|
def fake_row_to_model(*args, **kwargs) -> Model: # type: ignore[no-untyped-def]
|
|
return override_model
|
|
|
|
monkeypatch.setattr("routstr.payment.models._row_to_model", fake_row_to_model)
|
|
|
|
override_row = SimpleNamespace(id="azure/gpt-4o", upstream_provider_id=2, enabled=True)
|
|
|
|
_, provider_map, _ = create_model_mappings(
|
|
upstreams=[provider_a, provider_b],
|
|
overrides_by_key={("azure/gpt-4o", 2): (override_row, 1.01)},
|
|
disabled_model_keys=set(),
|
|
)
|
|
|
|
providers_for_alias = [p for _, p in provider_map["azure/gpt-4o"]]
|
|
assert provider_a in providers_for_alias
|
|
assert provider_b in providers_for_alias
|
|
assert len(providers_for_alias) == 2
|
|
|
|
|
|
def test_create_model_mappings_applies_override_only_to_matching_provider(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
"""Same-id overrides must not add provider-specific aliases to other providers."""
|
|
provider_a_model = create_test_model("same-id", prompt_price=0.01)
|
|
provider_a = create_test_provider(
|
|
"provider-a",
|
|
"https://provider-a.example/v1",
|
|
db_id=1,
|
|
models=[provider_a_model],
|
|
)
|
|
provider_b_model = create_test_model("same-id", prompt_price=0.02)
|
|
provider_b = create_test_provider(
|
|
"provider-b",
|
|
"https://provider-b.example/v1",
|
|
db_id=2,
|
|
models=[provider_b_model],
|
|
)
|
|
|
|
override_model = create_test_model("same-id", prompt_price=0.001)
|
|
override_model.alias_ids = ["provider-b-only"]
|
|
override_row = SimpleNamespace(id="same-id", upstream_provider_id=2, enabled=True)
|
|
|
|
def fake_row_to_model(*args, **kwargs) -> Model: # type: ignore[no-untyped-def]
|
|
return override_model
|
|
|
|
monkeypatch.setattr("routstr.payment.models._row_to_model", fake_row_to_model)
|
|
|
|
_, provider_map, _ = create_model_mappings(
|
|
upstreams=[provider_a, provider_b],
|
|
overrides_by_key={("same-id", 2): (override_row, 1.01)},
|
|
disabled_model_keys=set(),
|
|
)
|
|
|
|
assert [p for _, p in provider_map["provider-b-only"]] == [provider_b]
|
|
assert {p for _, p in provider_map["same-id"]} == {provider_a, provider_b}
|
|
|
|
|
|
def test_create_model_mappings_does_not_split_self_alias_from_base_identity(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
"""A write-time self-alias must not advertise one shared model twice (#639)."""
|
|
model_id = "deepseek/deepseek-chat"
|
|
provider_a = create_test_provider(
|
|
"provider-a",
|
|
"https://provider-a.example/v1",
|
|
db_id=1,
|
|
models=[create_test_model(model_id)],
|
|
)
|
|
provider_b = create_test_provider(
|
|
"provider-b",
|
|
"https://provider-b.example/v1",
|
|
db_id=2,
|
|
models=[create_test_model(model_id)],
|
|
)
|
|
|
|
admin_saved_model = create_test_model(model_id)
|
|
admin_saved_model.forwarded_model_id = model_id
|
|
override_row = SimpleNamespace(id=model_id, upstream_provider_id=1, enabled=True)
|
|
|
|
def fake_row_to_model(*args, **kwargs) -> Model: # type: ignore[no-untyped-def]
|
|
return admin_saved_model
|
|
|
|
monkeypatch.setattr("routstr.payment.models._row_to_model", fake_row_to_model)
|
|
|
|
_, _, unique_models = create_model_mappings(
|
|
upstreams=[provider_a, provider_b],
|
|
overrides_by_key={(model_id, 1): (override_row, 1.0)},
|
|
disabled_model_keys=set(),
|
|
)
|
|
|
|
advertised_ids = sorted(
|
|
model.forwarded_model_id or model.id for model in unique_models.values()
|
|
)
|
|
assert advertised_ids == ["deepseek-chat"]
|
|
|
|
|
|
def test_create_model_mappings_disables_only_matching_provider() -> None:
|
|
"""Disabled overrides are scoped to the provider row, not the shared model id."""
|
|
provider_a = create_test_provider(
|
|
"provider-a",
|
|
"https://provider-a.example/v1",
|
|
db_id=1,
|
|
models=[create_test_model("same-id")],
|
|
)
|
|
provider_b = create_test_provider(
|
|
"provider-b",
|
|
"https://provider-b.example/v1",
|
|
db_id=2,
|
|
models=[create_test_model("same-id")],
|
|
)
|
|
|
|
_, provider_map, _ = create_model_mappings(
|
|
upstreams=[provider_a, provider_b],
|
|
overrides_by_key={},
|
|
disabled_model_keys={("same-id", 2)},
|
|
)
|
|
|
|
assert [p for _, p in provider_map["same-id"]] == [provider_a]
|