122 lines
4.0 KiB
Python
122 lines
4.0 KiB
Python
"""Tests for the selector: ranking, filtering, capability gating, policy re-rank."""
|
|
from __future__ import annotations
|
|
|
|
import pytest
|
|
|
|
from cowork_local.core.routing.models import (
|
|
ModelAssessment,
|
|
ModelMetadata,
|
|
Policy,
|
|
ProbeResult,
|
|
TaskType,
|
|
)
|
|
from cowork_local.core.routing.selector import best_model, rank_models
|
|
|
|
|
|
def _assessment(
|
|
model_id,
|
|
*,
|
|
provider="anthropic",
|
|
quality=0.8,
|
|
latency_ms=500,
|
|
cost_in=0.001,
|
|
cost_out=0.003,
|
|
available=True,
|
|
caps=None,
|
|
task=TaskType.CODING,
|
|
probe_success=True,
|
|
) -> ModelAssessment:
|
|
meta = ModelMetadata(
|
|
provider=provider,
|
|
model_id=model_id,
|
|
cost_per_1k_input=cost_in,
|
|
cost_per_1k_output=cost_out,
|
|
max_context=100000,
|
|
capabilities=set(caps or []),
|
|
available=available,
|
|
)
|
|
probe = ProbeResult(
|
|
latency_ms=latency_ms, success=probe_success,
|
|
quality_score=quality, tokens_out=50,
|
|
)
|
|
return ModelAssessment(metadata=meta, probes={task.value: probe})
|
|
|
|
|
|
def test_empty_returns_no_best():
|
|
assert best_model([], TaskType.CODING) is None
|
|
|
|
|
|
def test_best_is_highest_quality_under_quality_policy():
|
|
weak = _assessment("weak", quality=0.3)
|
|
strong = _assessment("strong", quality=0.95)
|
|
best = best_model([weak, strong], TaskType.CODING, Policy.QUALITY)
|
|
assert best is not None
|
|
assert best.assessment.metadata.model_id == "strong"
|
|
|
|
|
|
def test_unavailable_excluded():
|
|
down = _assessment("down", quality=0.99, available=False)
|
|
up = _assessment("up", quality=0.5)
|
|
ranking = rank_models([down, up], TaskType.CODING)
|
|
keys = [c.assessment.metadata.model_id for c in ranking.ranked]
|
|
assert "down" not in keys
|
|
assert ranking.best.assessment.metadata.model_id == "up"
|
|
|
|
|
|
def test_failed_probe_excluded():
|
|
broken = _assessment("broken", quality=0.99, probe_success=False)
|
|
ok = _assessment("ok", quality=0.4)
|
|
best = best_model([broken, ok], TaskType.CODING)
|
|
assert best.assessment.metadata.model_id == "ok"
|
|
|
|
|
|
def test_missing_probe_for_task_excluded():
|
|
# Only has a CODING probe; asking for REASONING must exclude it.
|
|
coding_only = _assessment("c", task=TaskType.CODING)
|
|
assert best_model([coding_only], TaskType.REASONING) is None
|
|
|
|
|
|
def test_required_capability_filters_out_incapable():
|
|
no_vision = _assessment("text", quality=0.95, caps=[])
|
|
vision = _assessment("vision", quality=0.6, caps=["vision"])
|
|
best = best_model(
|
|
[no_vision, vision], TaskType.CODING, required_capabilities=["vision"]
|
|
)
|
|
assert best.assessment.metadata.model_id == "vision"
|
|
|
|
|
|
def test_policy_change_reranks_without_reprobe():
|
|
"""Same assessments, different policy → different winner, no re-probing."""
|
|
smart_pricey_slow = _assessment(
|
|
"opus", quality=0.95, latency_ms=6000, cost_in=0.015, cost_out=0.075
|
|
)
|
|
cheap_fast_ok = _assessment(
|
|
"haiku", quality=0.7, latency_ms=200, cost_in=0.0002, cost_out=0.0004
|
|
)
|
|
candidates = [smart_pricey_slow, cheap_fast_ok]
|
|
|
|
q_best = best_model(candidates, TaskType.CODING, Policy.QUALITY)
|
|
c_best = best_model(candidates, TaskType.CODING, Policy.COST)
|
|
l_best = best_model(candidates, TaskType.CODING, Policy.LATENCY)
|
|
|
|
assert q_best.assessment.metadata.model_id == "opus" # quality wins
|
|
assert c_best.assessment.metadata.model_id == "haiku" # cost wins
|
|
assert l_best.assessment.metadata.model_id == "haiku" # latency wins
|
|
|
|
|
|
def test_ranking_is_descending_and_stable():
|
|
a = _assessment("a", quality=0.9)
|
|
b = _assessment("b", quality=0.6)
|
|
c = _assessment("c", quality=0.3)
|
|
ranking = rank_models([b, c, a], TaskType.CODING, Policy.QUALITY)
|
|
scores = [rc.score for rc in ranking.ranked]
|
|
assert scores == sorted(scores, reverse=True)
|
|
assert [rc.assessment.metadata.model_id for rc in ranking.ranked] == ["a", "b", "c"]
|
|
|
|
|
|
def test_score_of_returns_zero_for_unranked():
|
|
a = _assessment("a", quality=0.9)
|
|
ranking = rank_models([a], TaskType.CODING)
|
|
assert ranking.score_of("anthropic/a") > 0
|
|
assert ranking.score_of("anthropic/missing") == 0.0
|