52 lines
2.3 KiB
Python
52 lines
2.3 KiB
Python
from __future__ import annotations
|
|
|
|
from agents.providers import DeterministicCodingProvider, DeterministicSolProvider
|
|
from control_plane.resources.models import ModelRequest, Resource, ResourceKind
|
|
from model_router.providers import SolProvider, providers_from_resources
|
|
from model_router.router import ModelCapability, ModelRequestContract, ModelRouter
|
|
|
|
|
|
def test_model_router_persists_sanitized_request_metadata() -> None:
|
|
Resource.objects.create(
|
|
name="Qwen Test",
|
|
kind=ResourceKind.MODEL,
|
|
provider="qwen",
|
|
roles=["CODING"],
|
|
)
|
|
router = ModelRouter({"qwen": DeterministicCodingProvider()}, persist_requests=True)
|
|
|
|
router.complete(ModelRequestContract(purpose=ModelCapability.CODING, prompt="secret-looking prompt must not persist"))
|
|
|
|
request = ModelRequest.objects.get()
|
|
assert request.status == "COMPLETE"
|
|
assert request.logical_role == "CODING"
|
|
assert request.request["prompt_chars"] > 0
|
|
assert request.request["contains_raw_prompt"] is False
|
|
assert "secret-looking" not in str(request.request)
|
|
|
|
|
|
def test_model_router_health_is_non_throwing() -> None:
|
|
router = ModelRouter({"qwen": DeterministicCodingProvider()})
|
|
|
|
assert router.health() == {"qwen": "AVAILABLE"}
|
|
|
|
|
|
def test_opencode_model_key_resources_load_as_distinct_providers() -> None:
|
|
for key in ["sol", "terra", "luna"]:
|
|
Resource.objects.create(name=key.title(), kind=ResourceKind.MODEL, provider="opencode", roles=["REASONING"], config={"model_key": key})
|
|
|
|
providers = providers_from_resources()
|
|
|
|
assert {"sol", "terra", "luna"}.issubset(providers)
|
|
assert all(isinstance(providers[key], SolProvider) for key in ["sol", "terra", "luna"])
|
|
|
|
|
|
def test_model_router_prefers_exact_model_key_resource_for_persisted_requests() -> None:
|
|
Resource.objects.create(name="Generic Opencode", kind=ResourceKind.MODEL, provider="opencode", roles=["REASONING"], config={})
|
|
Resource.objects.create(name="Terra", kind=ResourceKind.MODEL, provider="opencode", roles=["REASONING"], config={"model_key": "terra"})
|
|
router = ModelRouter({"terra": DeterministicSolProvider("{}")}, persist_requests=True)
|
|
|
|
router.complete(ModelRequestContract(purpose=ModelCapability.REASONING, model_hint="terra", prompt="{}"))
|
|
|
|
request = ModelRequest.objects.get()
|
|
assert request.model == "Terra"
|