wehub-resource-sync
444 lines
15 KiB
Python
444 lines
15 KiB
Python
"""Unit tests for the additive router budget gate.
|
|
|
|
The gate reads a session's accumulated spend (assembled by the router step
|
|
and passed in as plain data) and, when spend crosses a configured limit,
|
|
either warns (annotate + log, tier UNCHANGED) or caps (lower the tier). The
|
|
default state — no configured limit — is a complete no-op, and an unknown
|
|
spend/price SUSPENDS rather than acting on missing data.
|
|
|
|
Everything here runs on plain dataclasses: no classifier bundle, no
|
|
TurnContext, no network, no session storage. Spend figures are synthetic.
|
|
Old-vs-new parity of the full pipeline with the gate off is pinned separately
|
|
by ``test_routing_policy_parity.py``; the paired-run below re-proves it at the
|
|
engine level (default/None budget == a suspended budget == no gate).
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
|
|
from opensquilla.engine.routing import (
|
|
BudgetGateInput,
|
|
PolicyInputs,
|
|
RoutingDecision,
|
|
RoutingPolicyEngine,
|
|
apply_budget_gate,
|
|
budget_gate,
|
|
route_class_for_tier,
|
|
)
|
|
from opensquilla.engine.steps.squilla_router import _session_accumulated_spend
|
|
|
|
VALID_TIERS = ["c0", "c1", "c2", "c3"]
|
|
|
|
TIERS = {
|
|
"c0": {"model": "dummy-nano-1"},
|
|
"c1": {"model": "dummy-mini-1"},
|
|
"c2": {"model": "dummy-pro-1"},
|
|
"c3": {"model": "dummy-max-1"},
|
|
}
|
|
|
|
|
|
def budget_input(**overrides: object) -> BudgetGateInput:
|
|
knobs: dict[str, object] = {
|
|
"action": "warn",
|
|
"limit_usd": 1.0,
|
|
"spend_usd": 5.0,
|
|
"spend_source": "billed",
|
|
"session_key": "agent:budget:main",
|
|
}
|
|
knobs.update(overrides)
|
|
return BudgetGateInput(**knobs) # type: ignore[arg-type]
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# budget_gate — pure decision logic
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_off_when_input_none() -> None:
|
|
result = budget_gate("c2", valid_tiers=VALID_TIERS, budget=None)
|
|
assert (result.tier, result.outcome) == ("c2", "off")
|
|
|
|
|
|
def test_warn_over_limit_keeps_tier() -> None:
|
|
result = budget_gate("c2", valid_tiers=VALID_TIERS, budget=budget_input(spend_usd=5.0))
|
|
assert (result.tier, result.outcome, result.action) == ("c2", "warn", "warn")
|
|
assert result.spend_usd == 5.0
|
|
assert result.limit_usd == 1.0
|
|
|
|
|
|
def test_under_limit_is_noop() -> None:
|
|
result = budget_gate("c2", valid_tiers=VALID_TIERS, budget=budget_input(spend_usd=0.5))
|
|
assert (result.tier, result.outcome) == ("c2", "under_limit")
|
|
|
|
|
|
def test_at_limit_boundary_is_noop() -> None:
|
|
# spend == limit is not "exceeds": stay under.
|
|
result = budget_gate("c2", valid_tiers=VALID_TIERS, budget=budget_input(spend_usd=1.0))
|
|
assert result.outcome == "under_limit"
|
|
|
|
|
|
def test_cap_lowers_to_cap_tier() -> None:
|
|
result = budget_gate(
|
|
"c3",
|
|
valid_tiers=VALID_TIERS,
|
|
budget=budget_input(action="cap", cap_tier="c0", spend_usd=9.0),
|
|
)
|
|
assert (result.tier, result.outcome, result.from_tier) == ("c0", "cap", "c3")
|
|
|
|
|
|
def test_cap_never_raises_when_target_at_or_above_current() -> None:
|
|
# Working tier c0, cap target c2: capping would RAISE — degrade to warn.
|
|
result = budget_gate(
|
|
"c0",
|
|
valid_tiers=VALID_TIERS,
|
|
budget=budget_input(action="cap", cap_tier="c2", spend_usd=9.0),
|
|
)
|
|
assert (result.tier, result.outcome) == ("c0", "warn")
|
|
|
|
|
|
def test_cap_with_unknown_target_degrades_to_warn() -> None:
|
|
result = budget_gate(
|
|
"c3",
|
|
valid_tiers=VALID_TIERS,
|
|
budget=budget_input(action="cap", cap_tier="zz_unknown", spend_usd=9.0),
|
|
)
|
|
assert (result.tier, result.outcome) == ("c3", "warn")
|
|
|
|
|
|
def test_suspends_on_unknown_spend() -> None:
|
|
# No billed/estimated total yet (or a required price could not be priced):
|
|
# never act on missing data.
|
|
result = budget_gate(
|
|
"c2",
|
|
valid_tiers=VALID_TIERS,
|
|
budget=budget_input(spend_usd=None, spend_source="unknown"),
|
|
)
|
|
assert (result.tier, result.outcome) == ("c2", "suspended")
|
|
|
|
|
|
def test_cap_suspends_on_unknown_spend() -> None:
|
|
# A cap action must also suspend on unknown spend — no downgrade on a guess.
|
|
result = budget_gate(
|
|
"c3",
|
|
valid_tiers=VALID_TIERS,
|
|
budget=budget_input(action="cap", cap_tier="c0", spend_usd=None),
|
|
)
|
|
assert (result.tier, result.outcome) == ("c3", "suspended")
|
|
|
|
|
|
def test_forward_estimate_pushes_over_limit() -> None:
|
|
# spend alone (0.9) is under; spend + estimate (1.1) crosses.
|
|
result = budget_gate(
|
|
"c2",
|
|
valid_tiers=VALID_TIERS,
|
|
budget=budget_input(spend_usd=0.9, estimate_usd=0.2),
|
|
)
|
|
assert result.outcome == "warn"
|
|
assert result.projected_usd == pytest.approx(1.1)
|
|
|
|
|
|
def test_absent_estimate_compares_spend_alone() -> None:
|
|
result = budget_gate(
|
|
"c2",
|
|
valid_tiers=VALID_TIERS,
|
|
budget=budget_input(spend_usd=0.9, estimate_usd=None),
|
|
)
|
|
assert result.outcome == "under_limit"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# apply_budget_gate — decision + metadata effects
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _decision(tier: str = "c2") -> RoutingDecision:
|
|
return RoutingDecision(
|
|
tier=tier, model=TIERS[tier]["model"], confidence=0.9, source="v4_phase3"
|
|
)
|
|
|
|
|
|
def test_apply_off_writes_nothing() -> None:
|
|
extra: dict = {"route_class": "R2"}
|
|
meta: dict = {}
|
|
decision = _decision("c2")
|
|
result = budget_gate("c2", valid_tiers=VALID_TIERS, budget=None)
|
|
out = apply_budget_gate(decision, result, tiers=TIERS, extra=extra, metadata_updates=meta)
|
|
assert out is decision
|
|
assert meta == {}
|
|
assert extra == {"route_class": "R2"}
|
|
|
|
|
|
def test_apply_suspended_writes_nothing() -> None:
|
|
extra: dict = {"route_class": "R2"}
|
|
meta: dict = {}
|
|
decision = _decision("c2")
|
|
result = budget_gate(
|
|
"c2", valid_tiers=VALID_TIERS, budget=budget_input(spend_usd=None)
|
|
)
|
|
out = apply_budget_gate(decision, result, tiers=TIERS, extra=extra, metadata_updates=meta)
|
|
assert out.tier == "c2"
|
|
assert meta == {}
|
|
assert extra == {"route_class": "R2"}
|
|
|
|
|
|
def test_apply_under_limit_writes_nothing() -> None:
|
|
extra: dict = {"route_class": "R2"}
|
|
meta: dict = {}
|
|
decision = _decision("c2")
|
|
result = budget_gate(
|
|
"c2", valid_tiers=VALID_TIERS, budget=budget_input(spend_usd=0.1)
|
|
)
|
|
out = apply_budget_gate(decision, result, tiers=TIERS, extra=extra, metadata_updates=meta)
|
|
assert out.tier == "c2"
|
|
assert meta == {}
|
|
assert extra == {"route_class": "R2"}
|
|
|
|
|
|
def test_apply_warn_annotates_but_keeps_tier() -> None:
|
|
extra: dict = {"route_class": "R2"}
|
|
meta: dict = {}
|
|
decision = _decision("c2")
|
|
result = budget_gate("c2", valid_tiers=VALID_TIERS, budget=budget_input(spend_usd=5.0))
|
|
out = apply_budget_gate(decision, result, tiers=TIERS, extra=extra, metadata_updates=meta)
|
|
# Tier + model UNCHANGED.
|
|
assert (out.tier, out.model, out.source) == ("c2", "dummy-pro-1", "v4_phase3")
|
|
assert meta["router_budget_applied"] is True
|
|
assert meta["router_budget_outcome"] == "warn"
|
|
assert meta["router_budget_action"] == "warn"
|
|
assert meta["router_budget_spend_usd"] == 5.0
|
|
assert meta["router_budget_limit_usd"] == 1.0
|
|
trail = extra["routing_trail"]
|
|
assert trail[-1]["stage"] == "budget_gate"
|
|
assert trail[-1]["rule"] == "warn"
|
|
assert extra["budget_gate_applied"] is True
|
|
|
|
|
|
def test_apply_cap_rebinds_tier_and_model() -> None:
|
|
extra: dict = {"route_class": "R3", "final_tier": "c3"}
|
|
meta: dict = {}
|
|
decision = _decision("c3")
|
|
result = budget_gate(
|
|
"c3",
|
|
valid_tiers=VALID_TIERS,
|
|
budget=budget_input(action="cap", cap_tier="c1", spend_usd=9.0),
|
|
)
|
|
out = apply_budget_gate(decision, result, tiers=TIERS, extra=extra, metadata_updates=meta)
|
|
assert (out.tier, out.model, out.source) == ("c1", "dummy-mini-1", "budget_cap")
|
|
assert meta["router_budget_outcome"] == "cap"
|
|
assert meta["router_budget_from_tier"] == "c3"
|
|
assert meta["router_budget_to_tier"] == "c1"
|
|
assert extra["final_tier"] == "c1"
|
|
assert extra["final_route_class"] == route_class_for_tier("c1")
|
|
trail = extra["routing_trail"]
|
|
assert trail[-1] == {
|
|
"stage": "budget_gate",
|
|
"rule": "cap",
|
|
"spend_usd": 9.0,
|
|
"limit_usd": 1.0,
|
|
"spend_source": "billed",
|
|
"from_tier": "c3",
|
|
"to_tier": "c1",
|
|
}
|
|
|
|
|
|
def test_apply_warn_records_spend_source_in_trail_and_metadata() -> None:
|
|
# A warn outcome threads the spend basis into both the metadata trace and
|
|
# the routing-trail entry so the decision is auditable.
|
|
extra: dict = {"route_class": "R2"}
|
|
meta: dict = {}
|
|
decision = _decision("c2")
|
|
result = budget_gate(
|
|
"c2",
|
|
valid_tiers=VALID_TIERS,
|
|
budget=budget_input(spend_usd=5.0, spend_source="estimate_mixed"),
|
|
)
|
|
apply_budget_gate(decision, result, tiers=TIERS, extra=extra, metadata_updates=meta)
|
|
assert meta["router_budget_spend_source"] == "estimate_mixed"
|
|
assert extra["routing_trail"][-1]["spend_source"] == "estimate_mixed"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# _session_accumulated_spend — seeded provenance -> spend source
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _spend_ctx(**metadata: object) -> SimpleNamespace:
|
|
return SimpleNamespace(metadata=dict(metadata))
|
|
|
|
|
|
def test_accumulated_spend_billed_ignores_rollup_label() -> None:
|
|
# A positive billed total is authoritative regardless of the rollup label.
|
|
spend, source = _session_accumulated_spend(
|
|
_spend_ctx(
|
|
session_billed_cost_usd=3.0,
|
|
session_total_cost_usd=3.0,
|
|
session_cost_source="mixed",
|
|
)
|
|
)
|
|
assert (spend, source) == (3.0, "billed")
|
|
|
|
|
|
def test_accumulated_spend_estimate_label_maps_to_estimate() -> None:
|
|
spend, source = _session_accumulated_spend(
|
|
_spend_ctx(
|
|
session_billed_cost_usd=0.0,
|
|
session_total_cost_usd=2.5,
|
|
session_estimated_cost_usd=2.5,
|
|
session_cost_source="opensquilla_estimate",
|
|
)
|
|
)
|
|
assert (spend, source) == (2.5, "estimate")
|
|
|
|
|
|
def test_accumulated_spend_mixed_label_maps_to_estimate_mixed() -> None:
|
|
spend, source = _session_accumulated_spend(
|
|
_spend_ctx(
|
|
session_billed_cost_usd=0.0,
|
|
session_total_cost_usd=4.0,
|
|
session_cost_source="mixed",
|
|
)
|
|
)
|
|
assert (spend, source) == (4.0, "estimate_mixed")
|
|
|
|
|
|
def test_accumulated_spend_absent_label_stays_estimate() -> None:
|
|
spend, source = _session_accumulated_spend(_spend_ctx(session_total_cost_usd=1.5))
|
|
assert (spend, source) == (1.5, "estimate")
|
|
|
|
|
|
def test_accumulated_spend_unknown_and_none_preserved() -> None:
|
|
# No spend keys at all -> unknown (the gate suspends on missing data).
|
|
assert _session_accumulated_spend(_spend_ctx()) == (None, "unknown")
|
|
# Keys present but all zero -> a known-zero spend, not an unknown one, even
|
|
# with a rollup label present.
|
|
assert _session_accumulated_spend(
|
|
_spend_ctx(
|
|
session_billed_cost_usd=0.0,
|
|
session_total_cost_usd=0.0,
|
|
session_cost_source="mixed",
|
|
)
|
|
) == (0.0, "none")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Full-engine integration + default-path parity
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def router_cfg(**overrides: object) -> SimpleNamespace:
|
|
knobs: dict[str, object] = {
|
|
"default_tier": "c1",
|
|
"confidence_threshold": 0.5,
|
|
"confidence_high_tier_margin": 0.05,
|
|
"complaint_upgrade_enabled": True,
|
|
"complaint_upgrade_steps": 1,
|
|
"complaint_upgrade_max_chars": 160,
|
|
"kv_cache_anti_downgrade_enabled": True,
|
|
"kv_cache_anti_downgrade_window_seconds": 600,
|
|
}
|
|
knobs.update(overrides)
|
|
return SimpleNamespace(**knobs)
|
|
|
|
|
|
def make_inputs(
|
|
*,
|
|
tier: str = "c2",
|
|
confidence: float = 0.9,
|
|
tokens: int = 0,
|
|
budget: BudgetGateInput | None = None,
|
|
extra: dict | None = None,
|
|
) -> PolicyInputs:
|
|
return PolicyInputs(
|
|
decision=RoutingDecision(
|
|
tier=tier, model=TIERS[tier]["model"], confidence=confidence, source="v4_phase3"
|
|
),
|
|
message="please summarize this short note",
|
|
router_cfg=router_cfg(),
|
|
tiers=TIERS,
|
|
valid_tiers=VALID_TIERS,
|
|
routing_history=None,
|
|
extra=extra if extra is not None else {"route_class": route_class_for_tier(tier)},
|
|
thinking_mode="T2",
|
|
prompt_policy="P1",
|
|
history_strategy=True,
|
|
material_estimated_tokens=tokens,
|
|
context_window_tokens=200_000,
|
|
now=1000.0,
|
|
budget=budget,
|
|
)
|
|
|
|
|
|
def test_engine_warn_over_limit_keeps_routed_tier() -> None:
|
|
engine = RoutingPolicyEngine()
|
|
result = engine.run(
|
|
make_inputs(tier="c2", budget=budget_input(action="warn", limit_usd=1.0, spend_usd=9.0))
|
|
)
|
|
assert result.decision.tier == "c2" # warn never changes the tier
|
|
assert result.metadata_updates["router_budget_outcome"] == "warn"
|
|
assert result.metadata_updates["router_budget_applied"] is True
|
|
|
|
|
|
def test_engine_cap_lowers_routed_tier() -> None:
|
|
engine = RoutingPolicyEngine()
|
|
result = engine.run(
|
|
make_inputs(
|
|
tier="c3",
|
|
confidence=0.99,
|
|
budget=budget_input(action="cap", limit_usd=1.0, spend_usd=9.0, cap_tier="c0"),
|
|
)
|
|
)
|
|
assert result.decision.tier == "c0"
|
|
assert result.decision.model == "dummy-nano-1"
|
|
assert result.decision.source == "budget_cap"
|
|
assert result.metadata_updates["router_budget_to_tier"] == "c0"
|
|
|
|
|
|
def test_engine_default_budget_writes_no_budget_metadata() -> None:
|
|
engine = RoutingPolicyEngine()
|
|
result = engine.run(make_inputs(tier="c2", budget=None))
|
|
budget_keys = [k for k in result.metadata_updates if k.startswith("router_budget")]
|
|
assert budget_keys == []
|
|
|
|
|
|
# The parity matrix: default (None) budget must equal an active-but-suspended
|
|
# budget must equal an active-but-under-limit budget — all three are complete
|
|
# no-ops. This re-proves, at the engine level, that the additive stage does not
|
|
# perturb routing output when it does not act.
|
|
PARITY_SCENARIOS = [
|
|
{"tier": "c2", "confidence": 0.9, "tokens": 0},
|
|
{"tier": "c0", "confidence": 0.40, "tokens": 0},
|
|
{"tier": "c3", "confidence": 0.20, "tokens": 0},
|
|
{"tier": "c1", "confidence": 0.9, "tokens": 30_000},
|
|
{"tier": "c0", "confidence": 0.9, "tokens": 90_000},
|
|
]
|
|
|
|
|
|
def _engine_snapshot(scenario: dict, budget: BudgetGateInput | None) -> tuple:
|
|
engine = RoutingPolicyEngine()
|
|
extra: dict = {"route_class": route_class_for_tier(scenario["tier"])}
|
|
result = engine.run(make_inputs(budget=budget, extra=extra, **scenario))
|
|
return (
|
|
result.decision.tier,
|
|
result.decision.model,
|
|
result.decision.source,
|
|
result.decision.confidence,
|
|
result.thinking_mode,
|
|
result.prompt_policy,
|
|
extra,
|
|
result.metadata_updates,
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize("scenario", PARITY_SCENARIOS)
|
|
def test_default_budget_matches_suspended_and_under_limit(scenario: dict) -> None:
|
|
baseline = _engine_snapshot(scenario, None)
|
|
suspended = _engine_snapshot(
|
|
scenario, budget_input(limit_usd=1.0, spend_usd=None, spend_source="unknown")
|
|
)
|
|
under_limit = _engine_snapshot(scenario, budget_input(limit_usd=1_000_000.0, spend_usd=0.01))
|
|
assert baseline == suspended
|
|
assert baseline == under_limit
|