diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index e58c6d108ec..bcc90b9ebd5 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -59,7 +59,7 @@ from litellm.proxy.db.db_transaction_queue.tool_discovery_queue import ( from litellm.proxy.route_llm_request import ROUTE_ENDPOINT_MAPPING from litellm.proxy.spend_tracking.auto_router_session_queue import AutoRouterSessionQueue from litellm.proxy.spend_tracking.auto_router_sessions import ( - auto_router_group_kinds, + auto_router_kind, turn_from_spend_payload, ) from litellm.proxy.spend_tracking.compression_savings import ( @@ -303,12 +303,12 @@ class DBSpendUpdateWriter: session_id = payload.get("session_id") model_group = payload.get("model_group") model = payload.get("model") - start_time = payload.get("startTime") if prisma_client is None or llm_router is None or not session_id or not model_group or not model: return - router_kind = auto_router_group_kinds(llm_router).get(model_group) + router_kind = auto_router_kind(llm_router, model_group) if router_kind is None: return + start_time = payload.get("startTime") started_at = start_time if isinstance(start_time, datetime) else _parse_start_time(start_time) if started_at is None: return diff --git a/litellm/proxy/spend_tracking/auto_router_sessions.py b/litellm/proxy/spend_tracking/auto_router_sessions.py index c54cc068818..45f99817f1c 100644 --- a/litellm/proxy/spend_tracking/auto_router_sessions.py +++ b/litellm/proxy/spend_tracking/auto_router_sessions.py @@ -28,7 +28,7 @@ from pydantic import BaseModel, TypeAdapter, ValidationError import litellm from litellm._logging import verbose_proxy_logger -from litellm.router_utils.auto_router_model_naming import classify_strategy_router_model +from litellm.router_utils.auto_router_model_naming import StrategyRouterKind, classify_strategy_router_model if TYPE_CHECKING: from litellm.router import Router @@ -289,15 +289,44 @@ def _unclassified(turn: TurnFacts, baseline_spend: float, state: SessionState) - ) +def auto_router_kind(router: "Router", model_group: str) -> StrategyRouterKind | None: + """Which kind of auto-router serves ``model_group``, or None when none does. + + Answered from the registries the router already keys by public alias, so the + lookup costs the same on a proxy with a thousand models as on one with ten. + It runs on every request that carries a session, most of which never touched + an auto-router, so it must decide that before paying for anything else; a + scan of ``model_list`` charged the whole proxy for a feature few requests use. + + The alias is load-bearing rather than incidental: the auto-router's own + classifier sub-calls share the session but carry the judge model's group, so + keying on it folds one turn per routed request and no classifier noise. + + A complexity router with ``adaptive`` enabled owns an entry in both + registries, so complexity is answered first; that is the kind + ``auto_router_group_kinds`` reports for the same deployment, and the two have + to agree or the dashboard would label the group one way and the rows another. + """ + if model_group in router.complexity_routers: + return "complexity" + if model_group in router.quality_routers: + return "quality" + if model_group in router.adaptive_routers: + return "adaptive" + if model_group in router.auto_routers: + return "semantic" + return None + + def auto_router_group_kinds(router: "Router") -> Mapping[str, str]: """Public alias to router kind, for every auto-router on the proxy. ``model_name`` is what a caller sends and what spend rows record, while the ``litellm_params.model`` string carries the ``auto_router/...`` discriminator - that says it is one. Filtering turns by this mapping is the same filter the - dashboard has always used, and it is load-bearing: the auto-router's own - classifier sub-calls share the session but carry the judge model's group, so - keying on the alias yields one entry per routed turn with no classifier noise. + that says it is one. This is the read side's enumeration: which groups the + benchmarks dashboard asks the rollup about. The write path answers the same + question for one group at a time through ``auto_router_kind``, because a scan + of every deployment is fine once per dashboard load and not fine per request. Derived per call rather than cached because the router gains and loses deployments while it runs. diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index 191080e3a48..7e6c114f688 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -2236,3 +2236,129 @@ async def test_daily_transaction_compression_saved_tokens_zero_when_absent(): assert transaction["compression_saved_tokens"] == 0 assert transaction["compression_savings_spend"] == 0 assert transaction["prompt_caching_savings_spend"] == 0 + + +class _RecordingSessionQueue: + """Stands in for the session queue so a test sees exactly what one request staged.""" + + def __init__(self): + self.staged = [] + + async def record_turn(self, key, router_kind, baseline_model, turn) -> None: + self.staged.append((key, router_kind, turn.model)) + + +class _RegistryOnlyRouter: + """A router that can answer which of its groups an auto-router serves, and nothing else. + + ``model_list`` raises rather than returning: deciding that from an enumeration + of every deployment charged all traffic on the proxy, auto-routed or not, for a + fact the router already has keyed by group, and the cost grew with the number of + models configured. Reading it here fails the test instead of quietly costing + that again. + """ + + def __init__(self, complexity=(), quality=(), adaptive=(), semantic=()): + self.complexity_routers = dict.fromkeys(complexity, ()) + self.quality_routers = dict.fromkeys(quality, ()) + self.adaptive_routers = dict.fromkeys(adaptive, ()) + self.auto_routers = dict.fromkeys(semantic, ()) + + @property + def model_list(self): + raise AssertionError("recording a turn must not enumerate model_list") + + +def _turn_payload(model_group: str, metadata: str | None = None, session_id: str = "session-1") -> dict: + return { + "session_id": session_id, + "model_group": model_group, + "model": "anthropic/claude-haiku-4-5", + "model_id": "deployment-1", + "startTime": datetime(2026, 8, 4, 12, 0, tzinfo=timezone.utc), + "custom_llm_provider": "anthropic", + "spend": 0.01, + "prompt_tokens": 8000, + "completion_tokens": 200, + "total_tokens": 8200, + "metadata": json.dumps({"usage_object": {"cache_read_input_tokens": 6000}}) + if metadata is None + else metadata, + } + + +async def _record(router, payload: dict) -> _RecordingSessionQueue: + writer = DBSpendUpdateWriter() + queue = _RecordingSessionQueue() + writer.auto_router_session_queue = queue + with patch("litellm.proxy.proxy_server.llm_router", router): + await writer._record_auto_router_turn_unsafe(payload=payload, prisma_client=MagicMock()) + return queue + + +@pytest.mark.asyncio +class TestRecordingAnAutoRouterTurn: + async def test_a_turn_is_staged_without_ever_enumerating_the_model_list(self): + """What one turn costs must not grow with how many models the proxy serves.""" + router = _RegistryOnlyRouter(complexity=("smart-router",)) + + queue = await _record(router, _turn_payload("smart-router")) + + assert queue.staged == [(("session-1", "smart-router"), "complexity", "anthropic/claude-haiku-4-5")] + + @pytest.mark.parametrize( + "registry, kind", + [ + ("complexity", "complexity"), + ("quality", "quality"), + ("adaptive", "adaptive"), + ("semantic", "semantic"), + ], + ) + async def test_every_kind_of_auto_router_is_folded_under_its_own_kind(self, registry, kind): + """All four are auto-routers the benchmarks tab reports on, and only the three + strategy routers record a routing decision on the request; a semantic one + records none, so its traffic is recognised by the group it served.""" + router = _RegistryOnlyRouter(**{registry: ("a-router",)}) + + queue = await _record(router, _turn_payload("a-router")) + + assert [staged[1] for staged in queue.staged] == [kind] + + async def test_a_group_no_auto_router_serves_stages_nothing(self): + router = _RegistryOnlyRouter(complexity=("smart-router",)) + + queue = await _record(router, _turn_payload("gpt-4o")) + + assert queue.staged == [] + + async def test_the_routers_own_classifier_sub_call_is_not_a_turn(self): + """The classifier call shares the session with the request it classifies but + is billed to the judge model's group, so folding it would double every + auto-routed turn and price the session against a model nobody asked for.""" + router = _RegistryOnlyRouter(complexity=("smart-router",)) + classifier_call = _turn_payload( + "judge", + metadata=json.dumps({"usage_object": {}, "internal_call_origin": "autorouter_classifier"}), + ) + + queue = await _record(router, classifier_call) + + assert queue.staged == [] + + async def test_a_request_that_touched_no_auto_router_never_parses_its_metadata(self): + """Every spend write reaches here and almost none of them are auto-routed, so + the group is checked before anything is deserialized. Unparseable metadata + stands in for that parse: if it happens at all, this raises.""" + router = _RegistryOnlyRouter(complexity=("smart-router",)) + + queue = await _record(router, _turn_payload("gpt-4o", metadata="{not json")) + + assert queue.staged == [] + + async def test_a_request_without_a_session_is_not_a_turn(self): + router = _RegistryOnlyRouter(complexity=("smart-router",)) + + queue = await _record(router, _turn_payload("smart-router", metadata="{not json", session_id="")) + + assert queue.staged == [] diff --git a/tests/test_litellm/proxy/spend_tracking/test_auto_router_sessions.py b/tests/test_litellm/proxy/spend_tracking/test_auto_router_sessions.py index 995e4a137a0..9606ca2a176 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_auto_router_sessions.py +++ b/tests/test_litellm/proxy/spend_tracking/test_auto_router_sessions.py @@ -726,3 +726,103 @@ class TestFlushCostDoesNotGrowWithSessionCount: assert await queue.flush(_RecordingPrisma(table)) == 5 assert table.reads == 1 assert all(data["update"]["return_turns"] == {"increment": 1} for _, data in table.upserts) + + +def _router_with_auto_routers(): + """A real router carrying one auto-router of each kind that builds without extra packages. + + A semantic auto-router needs the ``semantic_router`` package to initialize, so + it is exercised against the registries directly in + ``TestResolvingOneGroupsKind`` rather than here. + """ + from litellm import Router + + return Router( + model_list=[ + { + "model_name": "adaptive-complexity-router", + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": { + "tiers": {"SIMPLE": ["cheap"], "MEDIUM": ["cheap"], "COMPLEX": ["pricey"]}, + "adaptive": True, + }, + "complexity_router_default_model": "cheap", + }, + }, + { + "model_name": "quality-router", + "litellm_params": { + "model": "auto_router/quality_router", + "quality_router_config": {"complexity_to_quality": {"SIMPLE": 1, "MEDIUM": 2, "COMPLEX": 3}}, + "quality_router_default_model": "cheap", + }, + }, + {"model_name": "cheap", "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-fake"}}, + {"model_name": "pricey", "litellm_params": {"model": "openai/gpt-4o", "api_key": "sk-fake"}}, + ] + ) + + +class _RegistriesOnly: + """Just the four registries a router keys its pre-routing strategies by.""" + + def __init__(self, complexity=(), quality=(), adaptive=(), semantic=()): + self.complexity_routers = dict.fromkeys(complexity, ()) + self.quality_routers = dict.fromkeys(quality, ()) + self.adaptive_routers = dict.fromkeys(adaptive, ()) + self.auto_routers = dict.fromkeys(semantic, ()) + + +class TestResolvingOneGroupsKind: + @pytest.mark.parametrize( + "registry, kind", + [ + ("complexity", "complexity"), + ("quality", "quality"), + ("adaptive", "adaptive"), + ("semantic", "semantic"), + ], + ) + def test_every_kind_of_auto_router_is_resolved_from_its_own_registry(self, registry, kind): + from litellm.proxy.spend_tracking.auto_router_sessions import auto_router_kind + + router = _RegistriesOnly(**{registry: ("a-router",)}) + assert auto_router_kind(router, "a-router") == kind + + def test_a_plain_model_group_resolves_to_nothing(self): + from litellm.proxy.spend_tracking.auto_router_sessions import auto_router_kind + + assert auto_router_kind(_RegistriesOnly(complexity=("a-router",)), "gpt-4o") is None + + +class TestTheWriteAndReadPathsAgreeOnWhatAnAutoRouterIs: + """The rollup is written per request and read per group, off two different + derivations of the same fact. A group the writer files under one kind and the + dashboard labels another, or filters out entirely, is a benchmark that reads + empty for traffic that really happened.""" + + def test_every_group_the_dashboard_asks_about_resolves_to_the_kind_it_labels(self): + from litellm.proxy.spend_tracking.auto_router_sessions import auto_router_group_kinds, auto_router_kind + + router = _router_with_auto_routers() + group_kinds = auto_router_group_kinds(router) + + assert dict(group_kinds) == {"adaptive-complexity-router": "complexity", "quality-router": "quality"} + assert {group: auto_router_kind(router, group) for group in group_kinds} == dict(group_kinds) + + def test_a_complexity_router_running_the_bandit_is_still_a_complexity_router(self): + """It owns an entry in both registries, so the lookup order decides, and only + one of the two orders agrees with what the dashboard labels the group.""" + router = _router_with_auto_routers() + + assert "adaptive-complexity-router" in router.complexity_routers + assert "adaptive-complexity-router" in router.adaptive_routers + + def test_a_group_no_auto_router_serves_is_left_out_of_both(self): + from litellm.proxy.spend_tracking.auto_router_sessions import auto_router_group_kinds, auto_router_kind + + router = _router_with_auto_routers() + + assert "cheap" not in auto_router_group_kinds(router) + assert auto_router_kind(router, "cheap") is None