mirror of
https://github.com/BerriAI/litellm.git
synced 2026-08-28 05:25:59 +00:00
feat(auto-router): track turns per complexity tier (LIT-5302) (#36209)
* feat(auto-router): track turns per complexity tier (LIT-5302)
Stamps complexity tier at decision time (rollup never re-derives from routed
model, since tier->model mapping is mutable config). Records per-tier turn
counts in LiteLLM_AutoRouterSession.tier_turns (jsonb), rolls up per router
in benchmarks SQL via jsonb_object_agg, returns on AutoRouterBenchmarkGroup
for dashboard turns/share metrics.
Addresses Greptile/Bugbot findings:
- Missing _SessionAggRow.tier_turns field: added with field_validator to
parse jsonb text cast and handle NULL. Would 500 every benchmarks read.
- Missing ::text cast on tier parameter: Postgres fails type inference on
parameterized CASE/IS NULL without explicit cast. Added to all usages.
- Docstring false claim (only complexity routers produce tiers): quality
router stamps numeric tier '1'/'2'/'3'. Per-type grouping in SQL prevents
cross-contamination. Rewrote docstring to clarify isolation.
- Comment convention violations: stripped per CLAUDE.md rule.
- Test gaps: 8 unit tests for extraction/validation/aggregation, 7 behavior
tests for SQL semantics against real Postgres. 12 mutations killed.
Fixed fragile complexity_router test that broke on nested function calls.
No API change; extends existing GET /auto_router/benchmarks response only.
Co-Authored-By: Claude <noreply@anthropic.com>
* fix(auto-router): address review findings on tier turns tracking
- Guard router_type update so a mid-session reconfigure can't pool
foreign tier names into tier_turns
- Keep pinned turns attributed to the tier that actually serves them
- Drop stray -- AlterTable comment from hand-written migration
- Drop the now-unnecessary ::text/json.loads round-trip; prisma
already returns tier_turns as a parsed dict
Co-Authored-By: Claude <noreply@anthropic.com>
* fix(auto-router): satisfy type-discipline lint gate
- tier_turns fields: dict[str, int] -> Mapping[str, int] (LIT001,
mutable collection in annotation); these are read-only after
construction
- _summed_agg_row: {} -> MappingProxyType({}) (LIT002, mutable dict
literal)
- default-fallback branch: replace the reassigned-without-Final
fallback_tier with a Final default_model_first flag and a single
ternary assignment (LIT010)
Verified locally: type_discipline_gate.py, ruff_strict_gate.py, and
type_check_gate.py all pass against the litellm_internal_staging
merge-base; full test_complexity_router.py (374), auto_router
management-endpoint tests (26), db-layer rollup tests (31), and the
live-Postgres proxy_behavior rollup suite (17) all pass.
Co-Authored-By: Claude <noreply@anthropic.com>
---------
Co-authored-by: Claude <noreply@anthropic.com>
This commit is contained in:
tin-berri
GitHub
parent
8db2fbaad0
commit
3238ce8406
12 changed files with 268 additions and 15 deletions
|
|
@ -0,0 +1 @@
|
|||
ALTER TABLE "LiteLLM_AutoRouterSession" ADD COLUMN IF NOT EXISTS "tier_turns" JSONB NOT NULL DEFAULT '{}';
|
||||
|
|
@ -1439,6 +1439,7 @@ model LiteLLM_AutoRouterSession {
|
|||
total_tokens BigInt @default(0)
|
||||
spend Float @default(0)
|
||||
saved_spend Float @default(0)
|
||||
tier_turns Json @default("{}")
|
||||
|
||||
@@id([api_key, session_id, router_name])
|
||||
@@index([last_turn_at], map: "idx_autorouter_session_last_turn")
|
||||
|
|
|
|||
|
|
@ -49,6 +49,7 @@ class AutoRouterTurnTransaction:
|
|||
cache_hit: bool
|
||||
cache_ttl_seconds: int | None
|
||||
cache_touched: bool
|
||||
tier: str | None = None
|
||||
|
||||
|
||||
class TurnCacheFacts(NamedTuple):
|
||||
|
|
@ -152,11 +153,13 @@ def build_autorouter_turn_transaction(
|
|||
return None
|
||||
usage_object_raw: Final = metadata.get("usage_object")
|
||||
cache: Final = turn_cache_facts(usage_object_raw if isinstance(usage_object_raw, Mapping) else None)
|
||||
tier_raw: Final = routing_decision.get("tier")
|
||||
return AutoRouterTurnTransaction(
|
||||
api_key=api_key,
|
||||
session_id=_bounded_session_id(session_id),
|
||||
router_name=router_name,
|
||||
router_type=str(routing_decision.get("router_type") or "unknown"),
|
||||
tier=tier_raw if isinstance(tier_raw, str) and tier_raw else None,
|
||||
model=model,
|
||||
turn_at=turn_at,
|
||||
total_tokens=int(payload.get("prompt_tokens") or 0) + int(payload.get("completion_tokens") or 0),
|
||||
|
|
@ -184,6 +187,8 @@ _COVERED: Final = _p("covered")
|
|||
_CACHE_HIT: Final = _p("cache_hit")
|
||||
_CACHE_TTL: Final = _p("cache_ttl_seconds")
|
||||
_TOUCHED: Final = _p("cache_touched")
|
||||
_TIER: Final = f"{_p('tier')}::text"
|
||||
_TIER_DELTA: Final = f"(CASE WHEN {_TIER} IS NULL THEN '{{}}'::jsonb ELSE jsonb_build_object({_TIER}, 1) END)"
|
||||
|
||||
_IN_ORDER: Final = f"{_TURN_AT}::timestamp >= t.last_turn_at"
|
||||
_SAME: Final = f"{_IN_ORDER} AND t.last_model = {_MODEL}"
|
||||
|
|
@ -201,7 +206,7 @@ INSERT INTO "LiteLLM_AutoRouterSession" AS t (
|
|||
last_model, models, turns, unordered_turns, covered_turns, cache_hits,
|
||||
same_model_turns, same_model_hits, first_visit_turns, first_visit_hits,
|
||||
return_turns, return_hits, return_expired_misses, return_within_ttl_misses,
|
||||
ttl_5m_turns, ttl_1h_turns, total_tokens, spend, saved_spend
|
||||
ttl_5m_turns, ttl_1h_turns, total_tokens, spend, saved_spend, tier_turns
|
||||
)
|
||||
VALUES (
|
||||
{_p("api_key")}, {_p("session_id")}, {_p("router_name")}, {_p("router_type")}, {_TURN_AT}::timestamp, {_TURN_AT}::timestamp,
|
||||
|
|
@ -211,7 +216,8 @@ VALUES (
|
|||
0, 0, 0, 0,
|
||||
(CASE WHEN {_CACHE_TTL}::int = {CACHE_TTL_5M_SECONDS} THEN 1 ELSE 0 END),
|
||||
(CASE WHEN {_CACHE_TTL}::int = {CACHE_TTL_1H_SECONDS} THEN 1 ELSE 0 END),
|
||||
{_p("total_tokens")}::bigint, {_p("spend")}::float8, {_p("saved_spend")}::float8
|
||||
{_p("total_tokens")}::bigint, {_p("spend")}::float8, {_p("saved_spend")}::float8,
|
||||
{_TIER_DELTA}
|
||||
)
|
||||
ON CONFLICT (api_key, session_id, router_name) DO UPDATE SET
|
||||
turns = t.turns + 1,
|
||||
|
|
@ -242,6 +248,9 @@ ON CONFLICT (api_key, session_id, router_name) DO UPDATE SET
|
|||
ELSE COALESCE((t.models -> {_MODEL} ->> 'ttl')::int, {_CACHE_TTL}::int) END)
|
||||
)),
|
||||
last_model = (CASE WHEN {_IN_ORDER} THEN {_MODEL} ELSE t.last_model END),
|
||||
tier_turns = (CASE WHEN {_TIER} IS NOT NULL AND t.router_type = {_p("router_type")}
|
||||
THEN t.tier_turns || jsonb_build_object({_TIER}, COALESCE((t.tier_turns ->> {_TIER})::int, 0) + 1)
|
||||
ELSE t.tier_turns END),
|
||||
first_turn_at = LEAST(t.first_turn_at, EXCLUDED.first_turn_at),
|
||||
last_turn_at = GREATEST(t.last_turn_at, EXCLUDED.last_turn_at)
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -4,8 +4,9 @@ AUTO ROUTER MANAGEMENT ENDPOINTS
|
|||
POST /auto_router/test_routing - Route one prompt through an unsaved complexity-router config
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
from collections.abc import Mapping, Sequence
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Annotated, Final
|
||||
|
||||
from pydantic import BaseModel, TypeAdapter
|
||||
|
|
@ -260,6 +261,7 @@ async def preview_auto_router_routing(
|
|||
class _SessionAggRow(BaseModel):
|
||||
router_name: str
|
||||
router_type: str
|
||||
tier_turns: Mapping[str, int]
|
||||
sessions: int
|
||||
turns: int
|
||||
unordered_turns: int
|
||||
|
|
@ -284,6 +286,23 @@ class _SessionAggRow(BaseModel):
|
|||
_SESSION_AGG_ROWS: Final = TypeAdapter(list[_SessionAggRow])
|
||||
|
||||
_BENCHMARKS_SQL: Final = """
|
||||
WITH windowed AS (
|
||||
SELECT * FROM "LiteLLM_AutoRouterSession"
|
||||
WHERE last_turn_at >= $1::timestamp AND first_turn_at < $2::timestamp
|
||||
),
|
||||
tier_maps AS (
|
||||
SELECT router_name, router_type, jsonb_object_agg(tier, tier_turns) AS tier_turns
|
||||
FROM (
|
||||
SELECT router_name, router_type, kv.key AS tier, SUM((kv.value)::int)::int AS tier_turns
|
||||
FROM windowed, LATERAL jsonb_each_text(tier_turns) AS kv
|
||||
GROUP BY router_name, router_type, kv.key
|
||||
) per_tier
|
||||
GROUP BY router_name, router_type
|
||||
)
|
||||
SELECT
|
||||
agg.*,
|
||||
COALESCE(tier_maps.tier_turns, '{}'::jsonb) AS tier_turns
|
||||
FROM (
|
||||
SELECT
|
||||
router_name,
|
||||
router_type,
|
||||
|
|
@ -306,10 +325,11 @@ SELECT
|
|||
COALESCE(SUM(spend), 0)::float8 AS spend,
|
||||
COALESCE(SUM(saved_spend), 0)::float8 AS saved_spend,
|
||||
COALESCE(SUM(EXTRACT(EPOCH FROM (last_turn_at - first_turn_at))), 0)::float8 AS session_seconds
|
||||
FROM "LiteLLM_AutoRouterSession"
|
||||
WHERE last_turn_at >= $1::timestamp AND first_turn_at < $2::timestamp
|
||||
FROM windowed
|
||||
GROUP BY router_name, router_type
|
||||
ORDER BY SUM(spend) DESC
|
||||
) agg
|
||||
LEFT JOIN tier_maps USING (router_name, router_type)
|
||||
ORDER BY agg.spend DESC
|
||||
"""
|
||||
|
||||
|
||||
|
|
@ -366,6 +386,7 @@ def _summed_agg_row(rows: Sequence[_SessionAggRow]) -> _SessionAggRow:
|
|||
return _SessionAggRow(
|
||||
router_name="",
|
||||
router_type="",
|
||||
tier_turns=MappingProxyType({}),
|
||||
sessions=sum(row.sessions for row in rows),
|
||||
turns=sum(row.turns for row in rows),
|
||||
unordered_turns=sum(row.unordered_turns for row in rows),
|
||||
|
|
@ -443,6 +464,7 @@ async def get_auto_router_benchmarks(
|
|||
AutoRouterBenchmarkGroup(
|
||||
router_name=row.router_name,
|
||||
router_type=row.router_type,
|
||||
tier_turns=row.tier_turns,
|
||||
**_benchmark_totals(row).model_dump(),
|
||||
)
|
||||
for row in rows
|
||||
|
|
|
|||
|
|
@ -1439,6 +1439,7 @@ model LiteLLM_AutoRouterSession {
|
|||
total_tokens BigInt @default(0)
|
||||
spend Float @default(0)
|
||||
saved_spend Float @default(0)
|
||||
tier_turns Json @default("{}")
|
||||
|
||||
@@id([api_key, session_id, router_name])
|
||||
@@index([last_turn_at], map: "idx_autorouter_session_last_turn")
|
||||
|
|
|
|||
|
|
@ -1730,6 +1730,7 @@ class ComplexityRouter(CustomLogger):
|
|||
routing_decision=self._build_routing_decision(
|
||||
routed_model=routed_model,
|
||||
cause=cause,
|
||||
tier=self._tier_for_model(routed_model),
|
||||
escalation_keyword=pin_escalation_keyword,
|
||||
escalated=escalated,
|
||||
conversation_continuing=conversation_continuing,
|
||||
|
|
@ -1797,7 +1798,8 @@ class ComplexityRouter(CustomLogger):
|
|||
|
||||
if user_message is None:
|
||||
verbose_router_logger.debug("ComplexityRouter: No user message found, routing to default model")
|
||||
if not self.config.plugins and self.config.default_model:
|
||||
default_model_first: Final = not self.config.plugins and self.config.default_model
|
||||
if default_model_first:
|
||||
# No plugins configured: preserve the pre-existing default_model-first
|
||||
# priority exactly (changing it would be a silent behavior change for
|
||||
# every non-plugin user, not just a security fix).
|
||||
|
|
@ -1809,12 +1811,14 @@ class ComplexityRouter(CustomLogger):
|
|||
routed_model = await self._pick_model_for_tier(
|
||||
ComplexityTier.MEDIUM, messages, resolved_messages, request_kwargs
|
||||
)
|
||||
fallback_tier: Final = None if default_model_first else ComplexityTier.MEDIUM
|
||||
return PreRoutingHookResponse(
|
||||
model=routed_model,
|
||||
messages=messages if has_original_messages else None,
|
||||
routing_decision=self._build_routing_decision(
|
||||
routed_model=routed_model,
|
||||
cause="default_fallback",
|
||||
tier=fallback_tier,
|
||||
conversation_continuing=conversation_continuing,
|
||||
),
|
||||
)
|
||||
|
|
|
|||
|
|
@ -2,6 +2,7 @@
|
|||
Types for auto-router management endpoints
|
||||
"""
|
||||
|
||||
from collections.abc import Mapping
|
||||
from typing import Final
|
||||
|
||||
from pydantic import BaseModel, Field, field_validator
|
||||
|
|
@ -120,6 +121,16 @@ class AutoRouterBenchmarkGroup(AutoRouterBenchmarkTotals):
|
|||
|
||||
router_name: str = Field(description="The auto-router alias requests were sent to")
|
||||
router_type: str = Field(description="complexity, adaptive or quality")
|
||||
tier_turns: Mapping[str, int] = Field(
|
||||
default_factory=dict,
|
||||
description="Turns per tier, keyed by the tier name the routing decision recorded at "
|
||||
"request time (never re-derived at read time, since the tier-to-model mapping is "
|
||||
"mutable config). Tier names are scoped to this group's router_type and are not "
|
||||
"comparable across types: a complexity router reports 'simple'/'medium'/'complex'/"
|
||||
"'reasoning', a quality router reports its numeric quality tier, and an adaptive router "
|
||||
"records no tier at all. Turns no tier served (the classifier fell back to default_model) "
|
||||
"are absent rather than pooled under a sentinel key, so the values may sum to less than turns",
|
||||
)
|
||||
|
||||
|
||||
class AutoRouterBenchmarksResponse(BaseModel):
|
||||
|
|
|
|||
|
|
@ -1439,6 +1439,7 @@ model LiteLLM_AutoRouterSession {
|
|||
total_tokens BigInt @default(0)
|
||||
spend Float @default(0)
|
||||
saved_spend Float @default(0)
|
||||
tier_turns Json @default("{}")
|
||||
|
||||
@@id([api_key, session_id, router_name])
|
||||
@@index([last_turn_at], map: "idx_autorouter_session_last_turn")
|
||||
|
|
|
|||
|
|
@ -38,11 +38,13 @@ async def _turn(
|
|||
tokens: int = 100,
|
||||
spend: float = 0.01,
|
||||
saved: float = 0.02,
|
||||
tier: "str | None" = None,
|
||||
) -> None:
|
||||
touched: Final = 1 if (hit or ttl is not None or not covered) else 0
|
||||
await db.execute_raw(
|
||||
UPSERT_AUTOROUTER_SESSION_SQL,
|
||||
key, session_id, router, router_type, model, at.isoformat(), tokens, spend, saved, covered, hit, ttl, touched,
|
||||
tier,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -195,6 +197,106 @@ async def test_a_reconfigured_alias_reports_each_router_type_as_its_own_group(db
|
|||
assert [(row["router_type"], row["sessions"]) for row in matching] == [("complexity", 1), ("quality", 1)]
|
||||
|
||||
|
||||
async def test_tier_turns_count_each_tier_that_served_a_turn(db):
|
||||
key = f"k-{uuid.uuid4()}"
|
||||
await _turn(db, key, "A", T0, tier="simple")
|
||||
await _turn(db, key, "B", T0 + timedelta(seconds=10), tier="complex")
|
||||
await _turn(db, key, "A", T0 + timedelta(seconds=20), tier="simple")
|
||||
|
||||
assert (await _row(db, key))["tier_turns"] == {"simple": 2, "complex": 1}
|
||||
|
||||
|
||||
async def test_an_untiered_turn_increments_no_tier_counter(db):
|
||||
key = f"k-{uuid.uuid4()}"
|
||||
await _turn(db, key, "A", T0, tier=None)
|
||||
assert (await _row(db, key))["tier_turns"] == {}
|
||||
|
||||
await _turn(db, key, "A", T0 + timedelta(seconds=10), tier="medium")
|
||||
await _turn(db, key, "A", T0 + timedelta(seconds=20), tier=None)
|
||||
row = await _row(db, key)
|
||||
assert row["tier_turns"] == {"medium": 1}
|
||||
assert row["turns"] == 3
|
||||
|
||||
|
||||
async def test_a_mid_session_router_type_change_keeps_foreign_tier_names_out_of_the_map(db):
|
||||
key = f"k-{uuid.uuid4()}"
|
||||
await _turn(db, key, "A", T0, router_type="complexity", tier="medium")
|
||||
await _turn(db, key, "A", T0 + timedelta(seconds=10), router_type="quality", tier="2")
|
||||
await _turn(db, key, "A", T0 + timedelta(seconds=20), router_type="complexity", tier="medium")
|
||||
|
||||
row = await _row(db, key)
|
||||
assert row["tier_turns"] == {"medium": 2}
|
||||
assert row["turns"] == 3
|
||||
|
||||
|
||||
async def test_an_out_of_order_turn_still_counts_toward_its_tier(db):
|
||||
key = f"k-{uuid.uuid4()}"
|
||||
await _turn(db, key, "A", T0 + timedelta(seconds=60), tier="simple")
|
||||
await _turn(db, key, "A", T0, tier="simple")
|
||||
|
||||
row = await _row(db, key)
|
||||
assert row["tier_turns"] == {"simple": 2}
|
||||
assert row["unordered_turns"] == 1
|
||||
|
||||
|
||||
async def test_the_benchmarks_aggregate_sums_tier_turns_across_sessions(db):
|
||||
key = f"k-{uuid.uuid4()}"
|
||||
router = f"r-{uuid.uuid4()}"
|
||||
await _turn(db, key, "A", T0, session_id=f"s-{uuid.uuid4()}", router=router, tier="simple")
|
||||
await _turn(db, key, "A", T0 + timedelta(seconds=10), session_id=f"s-{uuid.uuid4()}", router=router, tier="simple")
|
||||
await _turn(db, key, "B", T0 + timedelta(seconds=20), session_id=f"s-{uuid.uuid4()}", router=router, tier="complex")
|
||||
await _turn(db, key, "C", T0 + timedelta(seconds=30), session_id=f"s-{uuid.uuid4()}", router=router, tier=None)
|
||||
|
||||
rows = await db.query_raw(
|
||||
_BENCHMARKS_SQL,
|
||||
(T0 - timedelta(days=1)).isoformat(),
|
||||
(T0 + timedelta(days=1)).isoformat(),
|
||||
)
|
||||
grouped = next(row for row in rows if row["router_name"] == router)
|
||||
assert grouped["tier_turns"] == {"simple": 2, "complex": 1}
|
||||
assert grouped["turns"] == 4
|
||||
|
||||
|
||||
async def test_tier_maps_stay_separate_per_router_type_on_a_reconfigured_alias(db):
|
||||
key = f"k-{uuid.uuid4()}"
|
||||
router = f"r-{uuid.uuid4()}"
|
||||
await _turn(
|
||||
db, key, "A", T0, session_id=f"s-{uuid.uuid4()}", router=router, router_type="complexity", tier="medium"
|
||||
)
|
||||
await _turn(
|
||||
db,
|
||||
key,
|
||||
"A",
|
||||
T0 + timedelta(seconds=10),
|
||||
session_id=f"s-{uuid.uuid4()}",
|
||||
router=router,
|
||||
router_type="quality",
|
||||
tier="2",
|
||||
)
|
||||
|
||||
rows = await db.query_raw(
|
||||
_BENCHMARKS_SQL,
|
||||
(T0 - timedelta(days=1)).isoformat(),
|
||||
(T0 + timedelta(days=1)).isoformat(),
|
||||
)
|
||||
by_type = {row["router_type"]: row["tier_turns"] for row in rows if row["router_name"] == router}
|
||||
assert by_type == {"complexity": {"medium": 1}, "quality": {"2": 1}}
|
||||
|
||||
|
||||
async def test_a_window_with_no_tiered_turns_aggregates_to_an_empty_map(db):
|
||||
key = f"k-{uuid.uuid4()}"
|
||||
router = f"r-{uuid.uuid4()}"
|
||||
await _turn(db, key, "A", T0, session_id=f"s-{uuid.uuid4()}", router=router, tier=None)
|
||||
|
||||
rows = await db.query_raw(
|
||||
_BENCHMARKS_SQL,
|
||||
(T0 - timedelta(days=1)).isoformat(),
|
||||
(T0 + timedelta(days=1)).isoformat(),
|
||||
)
|
||||
grouped = next(row for row in rows if row["router_name"] == router)
|
||||
assert grouped["tier_turns"] == {}
|
||||
|
||||
|
||||
async def test_a_miss_that_touched_no_cache_does_not_advance_the_ttl_clock(db):
|
||||
key = f"k-{uuid.uuid4()}"
|
||||
await _turn(db, key, "A", T0, ttl=300)
|
||||
|
|
|
|||
|
|
@ -93,6 +93,19 @@ class TestBuildTransaction:
|
|||
def test_requests_without_a_routing_decision_are_skipped(self, metadata: dict):
|
||||
assert _build(metadata=metadata) is None
|
||||
|
||||
def test_the_tier_the_decision_recorded_is_carried_onto_the_transaction(self):
|
||||
transaction = _build(metadata=_metadata(routing_decision={**ROUTING_DECISION, "tier": "reasoning"}))
|
||||
assert transaction is not None and transaction.tier == "reasoning"
|
||||
|
||||
@pytest.mark.parametrize("tier", [None, "", 3, {"tier": "medium"}])
|
||||
def test_a_decision_without_a_usable_tier_records_no_tier(self, tier: object):
|
||||
transaction = _build(metadata=_metadata(routing_decision={**ROUTING_DECISION, "tier": tier}))
|
||||
assert transaction is not None and transaction.tier is None
|
||||
|
||||
def test_a_decision_that_never_mentions_tier_records_no_tier(self):
|
||||
transaction = _build()
|
||||
assert transaction is not None and transaction.tier is None
|
||||
|
||||
def test_router_name_falls_back_to_the_payload_model_group(self):
|
||||
transaction = _build(metadata=_metadata(routing_decision={"router_type": "complexity"}))
|
||||
assert transaction is not None and transaction.router_name == "live-auto"
|
||||
|
|
@ -167,7 +180,11 @@ class _FakeClient:
|
|||
self.db = _FakeDB(failures, poison_session)
|
||||
|
||||
|
||||
def _transaction(session_id: str = "s1", at: datetime = datetime(2026, 8, 1, 12, 0, 0)) -> AutoRouterTurnTransaction:
|
||||
def _transaction(
|
||||
session_id: str = "s1",
|
||||
at: datetime = datetime(2026, 8, 1, 12, 0, 0),
|
||||
tier: str | None = "medium",
|
||||
) -> AutoRouterTurnTransaction:
|
||||
return AutoRouterTurnTransaction(
|
||||
api_key="k1",
|
||||
session_id=session_id,
|
||||
|
|
@ -182,6 +199,7 @@ def _transaction(session_id: str = "s1", at: datetime = datetime(2026, 8, 1, 12,
|
|||
cache_hit=False,
|
||||
cache_ttl_seconds=None,
|
||||
cache_touched=False,
|
||||
tier=tier,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -201,7 +219,7 @@ class TestFlush:
|
|||
assert sql == UPSERT_AUTOROUTER_SESSION_SQL
|
||||
assert params == (
|
||||
"k1", "s1", "live-auto", "complexity", "bedrock/haiku",
|
||||
"2026-08-01T12:00:00", 100, 0.01, 0.02, 1, 0, None, 0,
|
||||
"2026-08-01T12:00:00", 100, 0.01, 0.02, 1, 0, None, 0, "medium",
|
||||
)
|
||||
|
||||
def test_a_connect_error_retries_the_same_statement(self):
|
||||
|
|
|
|||
|
|
@ -292,6 +292,7 @@ class TestAutoRouterBenchmarks:
|
|||
ROW = _SessionAggRow(
|
||||
router_name="live-auto",
|
||||
router_type="complexity",
|
||||
tier_turns={},
|
||||
sessions=4,
|
||||
turns=40,
|
||||
unordered_turns=1,
|
||||
|
|
@ -377,6 +378,21 @@ class TestAutoRouterBenchmarks:
|
|||
assert totals.avg_turns_per_session == 10.0
|
||||
assert totals.spend == 10.0
|
||||
|
||||
def test_tier_names_stay_scoped_to_the_router_type_that_recorded_them(self):
|
||||
quality = self.ROW.model_copy(
|
||||
update={"router_name": "quality-auto", "router_type": "quality", "tier_turns": {"2": 7}}
|
||||
)
|
||||
complexity = self.ROW.model_copy(update={"tier_turns": {"medium": 7}})
|
||||
assert complexity.tier_turns == {"medium": 7}
|
||||
assert quality.tier_turns == {"2": 7}
|
||||
|
||||
def test_summed_totals_carry_no_tier_map_because_names_are_router_scoped(self):
|
||||
from litellm.proxy.management_endpoints.auto_router_endpoints import _summed_agg_row
|
||||
|
||||
quality = self.ROW.model_copy(update={"router_type": "quality", "tier_turns": {"2": 7}})
|
||||
complexity = self.ROW.model_copy(update={"tier_turns": {"medium": 7}})
|
||||
assert _summed_agg_row([complexity, quality]).tier_turns == {}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_non_admin_roles_cannot_read_benchmarks(self):
|
||||
from litellm.proxy.management_endpoints.auto_router_endpoints import get_auto_router_benchmarks
|
||||
|
|
@ -427,3 +443,26 @@ class TestAutoRouterBenchmarks:
|
|||
assert response.routers_in_scope == 1
|
||||
assert response.groups[0].router_name == "live-auto"
|
||||
assert response.groups[0].saved_pct == response.totals.saved_pct == 75.0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"wire_value, expected", [({"simple": 24, "complex": 16}, {"simple": 24, "complex": 16}), ({}, {})]
|
||||
)
|
||||
async def test_the_tier_map_reaches_the_response_as_the_jsonb_column_returns_it(
|
||||
self, wire_value: dict, expected: dict, monkeypatch: pytest.MonkeyPatch
|
||||
):
|
||||
from litellm.proxy import proxy_server
|
||||
from litellm.proxy.management_endpoints.auto_router_endpoints import get_auto_router_benchmarks
|
||||
|
||||
class _DB:
|
||||
async def query_raw(self, sql: str, *params: object):
|
||||
return [{**TestAutoRouterBenchmarks.ROW.model_dump(), "tier_turns": wire_value}]
|
||||
|
||||
monkeypatch.setattr(proxy_server, "prisma_client", type("P", (), {"db": _DB()})())
|
||||
|
||||
response = await get_auto_router_benchmarks(
|
||||
user_api_key_dict=ADMIN,
|
||||
start_date="2026-07-01",
|
||||
end_date="2026-08-01",
|
||||
)
|
||||
assert response.groups[0].tier_turns == expected
|
||||
|
|
|
|||
|
|
@ -3477,6 +3477,23 @@ class TestSessionAffinity:
|
|||
# Pinned to the first turn's model, not re-classified down to SIMPLE.
|
||||
assert second.model == "o1-preview"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_pinned_turn_reports_the_tier_that_serves_it(self, mock_router_instance, session_affinity_config):
|
||||
mock_router_instance.cache = DualCache()
|
||||
router = ComplexityRouter(
|
||||
model_name="test-router",
|
||||
litellm_router_instance=mock_router_instance,
|
||||
complexity_router_config=session_affinity_config,
|
||||
)
|
||||
request_kwargs = self._request_kwargs("session-1")
|
||||
await router.async_pre_routing_hook(
|
||||
model="test-model", request_kwargs=request_kwargs, messages=self.REASONING_MESSAGE
|
||||
)
|
||||
pinned = await router.async_pre_routing_hook(
|
||||
model="test-model", request_kwargs=request_kwargs, messages=self.SIMPLE_MESSAGE
|
||||
)
|
||||
assert pinned.routing_decision["tier"] == "REASONING"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_different_sessions_classify_independently(self, mock_router_instance, session_affinity_config):
|
||||
mock_router_instance.cache = DualCache()
|
||||
|
|
@ -4362,7 +4379,24 @@ class TestRoutingDecisionContents:
|
|||
assert decision is not None
|
||||
assert decision["cause"] == "default_fallback"
|
||||
assert decision["routed_model"] == response.model
|
||||
assert "tier" not in decision
|
||||
assert decision.get("tier") == "MEDIUM"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_default_model_fallback_claims_no_tier(self, mock_router_instance, basic_config):
|
||||
router = ComplexityRouter(
|
||||
model_name="test-complexity-router",
|
||||
litellm_router_instance=mock_router_instance,
|
||||
complexity_router_config={**basic_config, "default_model": "gpt-4o"},
|
||||
)
|
||||
response = await router.async_pre_routing_hook(
|
||||
model="test-complexity-router",
|
||||
request_kwargs={},
|
||||
messages=[{"role": "system", "content": "be nice"}],
|
||||
)
|
||||
assert response is not None
|
||||
assert response.routing_decision is not None
|
||||
assert response.routing_decision["cause"] == "default_fallback"
|
||||
assert "tier" not in response.routing_decision
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_session_pin_decision(self, mock_router_instance, basic_config):
|
||||
|
|
@ -5907,11 +5941,21 @@ class TestConversationShapeDiscriminator:
|
|||
)
|
||||
builds = source.split("self._build_routing_decision(")[1:]
|
||||
assert builds
|
||||
missing = [
|
||||
i
|
||||
for i, block in enumerate(builds)
|
||||
if "conversation_continuing=conversation_continuing" not in block.split("),")[0]
|
||||
]
|
||||
missing = []
|
||||
for i, block in enumerate(builds):
|
||||
depth = 0
|
||||
end = 0
|
||||
for j, char in enumerate(block):
|
||||
if char == "(":
|
||||
depth += 1
|
||||
elif char == ")":
|
||||
depth -= 1
|
||||
if depth < 0:
|
||||
end = j
|
||||
break
|
||||
extracted = block[:end]
|
||||
if "conversation_continuing=conversation_continuing" not in extracted:
|
||||
missing.append(i)
|
||||
assert not missing, f"routing decisions {missing} do not carry the conversation shape"
|
||||
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue