diff --git a/litellm/proxy/db/autorouter_quality_signals.py b/litellm/proxy/db/autorouter_quality_signals.py index d5467d5f3c6..268e72b6a0e 100644 --- a/litellm/proxy/db/autorouter_quality_signals.py +++ b/litellm/proxy/db/autorouter_quality_signals.py @@ -39,13 +39,7 @@ from litellm._logging import verbose_proxy_logger if TYPE_CHECKING: from litellm.router import Router -# A cohort smaller than this is noise: one unlucky session moves the rate by whole points, -# and a comparison drawn from it would read as fact. MIN_COHORT_SESSIONS: Final = 20 - -# Below this share of turns carrying a caller-supplied session id, "sessions" in the -# non-routed cohort are mostly one-request artefacts of the fallback uuid, and any -# within-session signal computed over them is measuring the fallback, not the traffic. MIN_SESSION_ID_COVERAGE: Final = 0.8 @@ -56,12 +50,21 @@ class Turn: Kept as a plain object rather than a row dict so the signal functions state their inputs. """ - __slots__ = ("session_id", "model", "started_at", "client_disconnected", "router_name", "has_client_session_id") + __slots__ = ( + "session_id", + "api_key", + "model", + "started_at", + "client_disconnected", + "router_name", + "has_client_session_id", + ) def __init__( self, *, session_id: str, + api_key: str, model: str, started_at: float, client_disconnected: bool, @@ -69,6 +72,7 @@ class Turn: has_client_session_id: bool, ) -> None: self.session_id = session_id + self.api_key = api_key self.model = model self.started_at = started_at self.client_disconnected = client_disconnected @@ -113,7 +117,8 @@ def rank_models_by_cost(router: "Router", models: Iterable[str]) -> Mapping[str, if quote is None: verbose_proxy_logger.debug("quality signals: cannot price %s, leaving it unranked", model) priced: Final = sorted((quote[0], model) for model, quote in quotes if quote is not None) - return MappingProxyType({model: rank for rank, (_, model) in enumerate(priced)}) + distinct_costs: Final = tuple(dict.fromkeys(cost for cost, _ in priced)) + return MappingProxyType({model: distinct_costs.index(cost) for cost, model in priced}) def session_escalated(turns: Sequence[Turn], ranks: Mapping[str, int]) -> bool: @@ -154,28 +159,32 @@ def could_escalate(turns: Sequence[Turn], ranks: Mapping[str, int], reachable: I def _group_by_session(turns: Iterable[Turn]) -> Mapping[str, tuple[Turn, ...]]: - ordered: Final = sorted(turns, key=lambda turn: turn.session_id) + ordered: Final = sorted(turns, key=lambda turn: (turn.api_key, turn.session_id)) return MappingProxyType( - {session_id: tuple(group) for session_id, group in groupby(ordered, key=lambda turn: turn.session_id)} + {key: tuple(group) for key, group in groupby(ordered, key=lambda turn: (turn.api_key, turn.session_id))} ) def signals_for_cohort( turns: Sequence[Turn], ranks: Mapping[str, int], - reachable: Iterable[str], + reachable_by_key: Mapping[str, Iterable[str]], ) -> CohortSignals: """Both rates over the sessions that could have escalated. Abandonment shares escalation's denominator on purpose. The two numbers sit next to each other in the UI and get read as one population; computing them over different sets would make that reading wrong in a way nothing on screen would reveal. + + Reachability is looked up per session's own api_key, not pooled across every key in the + cohort: keys can carry different model lists, and a shared pool would make a key that + cannot reach a costlier model borrow one it never had, inventing escalation opportunities + that were never actually available to it. """ - reachable_models: Final = tuple(reachable) eligible: Final = tuple( session_turns for session_turns in _group_by_session(turns).values() - if could_escalate(session_turns, ranks, reachable_models) + if could_escalate(session_turns, ranks, reachable_by_key.get(session_turns[0].api_key, ())) ) if not eligible: return CohortSignals(sessions=0, escalation_rate_pct=None, abandonment_rate_pct=None) @@ -195,7 +204,11 @@ def baseline_unavailable_reason(turns: Sequence[Turn], cohort: CohortSignals) -> Session-id coverage is checked before size because the two failures need different words: a deployment that never sends session ids has plenty of rows and no sessions, and telling - it "not enough traffic" would send it looking for volume it already has. + it "not enough traffic" would send it looking for volume it already has. Below + ``MIN_SESSION_ID_COVERAGE``, "sessions" in the non-routed cohort are mostly one-request + artefacts of the fallback uuid, and any within-session signal computed over them is + measuring the fallback, not the traffic. Below ``MIN_COHORT_SESSIONS``, one unlucky session + moves the rate by whole points and a comparison drawn from it would read as fact. """ if turns: with_client_id: Final = sum(1 for turn in turns if turn.has_client_session_id) diff --git a/litellm/proxy/management_endpoints/auto_router_endpoints.py b/litellm/proxy/management_endpoints/auto_router_endpoints.py index b24c5ac298a..39c84b9c566 100644 --- a/litellm/proxy/management_endpoints/auto_router_endpoints.py +++ b/litellm/proxy/management_endpoints/auto_router_endpoints.py @@ -495,21 +495,6 @@ class _QualityTurnRow(BaseModel): _QUALITY_TURN_ROWS: Final = TypeAdapter(list[_QualityTurnRow]) -# Both cohorts come from this one scan, so they cannot drift apart: same window, same -# session grouping, same disconnect test. `router_name` is NULL for a directly-addressed -# request, which is what separates the two populations downstream. -# -# Abandonment reads error_information.error_code rather than `status`, because a client -# disconnect is still billed as a success -- partial streamed spend is deliberately charged -# on disconnect -- so `status` cannot see it. -# -# session_turn_count is how many rows in the window share this session_id. There is no -# persisted flag for "the caller supplied this id rather than the spend writer minting a -# fallback uuid" (`_get_session_id_for_spend_log`), but a minted fallback is always unique to -# its one request, so a session_id that repeats could only have come from the caller. Counting -# it this way needs nothing beyond columns every deployment already writes, prompt storage on -# or off, which a check against proxy_server_request (only persisted when prompt storage is on) -# would not have been. _QUALITY_SIGNALS_SQL: Final = """ SELECT session_id, @@ -517,7 +502,7 @@ SELECT EXTRACT(EPOCH FROM "startTime")::float8 AS started_at, (metadata #>> '{routing_decision,router_model_name}') AS router_name, COALESCE(metadata #>> '{error_information,error_code}' = '499', false) AS client_disconnected, - COUNT(*) OVER (PARTITION BY session_id)::int AS session_turn_count, + COUNT(*) OVER (PARTITION BY api_key, session_id)::int AS session_turn_count, api_key FROM "LiteLLM_SpendLogs" WHERE "startTime" >= $1::timestamp AND "startTime" < $2::timestamp @@ -539,11 +524,14 @@ def _quality_signals_for( restricting to keys that actually used the router keeps both sides of the comparison inside the same reachable-model world. - Reachability is judged against every model the key exercised anywhere in the window, not - against the models this one cohort happened to use. Scoping it to the cohort makes a - cohort that never escalated define its own ceiling, so every one of its sessions looks - like it had nowhere to go and the rate reads as unmeasurable rather than as zero -- - silently deleting exactly the well-behaved traffic the comparison exists to show. + Each session's eligibility is judged against every model its own api_key exercised + anywhere in the window, not against the models this one cohort happened to use or a pool + shared across keys. Scoping it to the cohort would make a cohort that never escalated + define its own ceiling, so every one of its sessions looks like it had nowhere to go and + the rate reads as unmeasurable rather than as zero -- silently deleting exactly the + well-behaved traffic the comparison exists to show. Pooling reachability across keys would + let a key that cannot reach a costlier model borrow one it never had, from another key + that only happens to share this cohort. """ from litellm.proxy.db.autorouter_quality_signals import ( Turn, @@ -558,26 +546,32 @@ def _quality_signals_for( router_keys: Final = frozenset(row.api_key for row in routed_rows) baseline_rows: Final = tuple(row for row in turns if row.router_name is None and row.api_key in router_keys) - reachable: Final = tuple(frozenset(row.model for row in turns if row.api_key in router_keys)) - ranks: Final = rank_models_by_cost(llm_router, reachable) if llm_router is not None else MappingProxyType({}) + router_key_rows: Final = tuple(row for row in turns if row.api_key in router_keys) + reachable_by_key: Final = MappingProxyType( + { + key: frozenset(row.model for row in router_key_rows if row.api_key == key) + for key in frozenset(row.api_key for row in router_key_rows) + } + ) + reachable_models: Final = tuple(frozenset(row.model for row in router_key_rows)) + ranks: Final = rank_models_by_cost(llm_router, reachable_models) if llm_router is not None else MappingProxyType({}) def as_turns(rows: Sequence[_QualityTurnRow]) -> tuple[Turn, ...]: return tuple( Turn( session_id=row.session_id, + api_key=row.api_key, model=row.model, started_at=row.started_at, client_disconnected=row.client_disconnected, router_name=row.router_name, - # A fallback uuid (_get_session_id_for_spend_log) is unique to its one - # request; only a caller-supplied session id repeats across rows. has_client_session_id=row.session_turn_count > 1, ) for row in rows ) - routed: Final = signals_for_cohort(as_turns(routed_rows), ranks, reachable) - baseline: Final = signals_for_cohort(as_turns(baseline_rows), ranks, reachable) + routed: Final = signals_for_cohort(as_turns(routed_rows), ranks, reachable_by_key) + baseline: Final = signals_for_cohort(as_turns(baseline_rows), ranks, reachable_by_key) unavailable: Final = baseline_unavailable_reason(as_turns(baseline_rows), baseline) return AutoRouterQualitySignals( router_name=router_name, @@ -614,8 +608,17 @@ async def get_auto_router_quality_signals( model mid-session, and how often they hung up mid-stream, for auto-routed traffic and for the same keys' directly-addressed traffic. - Reads LiteLLM_SpendLogs rather than the per-session rollup, because escalation is a - question about turn order and the rollup folds order away. + Both cohorts come from one scan of LiteLLM_SpendLogs rather than the per-session rollup, + so they cannot drift apart -- same window, same session grouping, same disconnect test -- + and because escalation is a question about turn order, which the rollup folds away. + `router_name` is NULL for a directly-addressed request, which is what separates the two + populations downstream. Abandonment reads `error_information.error_code` rather than + `status`, because a client disconnect still bills its partial streamed spend as a success + and so does not show up in `status`. `session_turn_count` counts, per api_key, how many + rows in the window share a session_id; a fallback uuid minted by the spend writer + (`_get_session_id_for_spend_log`) is always unique to its one request, so a repeating + session_id can only have come from the caller, which needs nothing beyond columns every + deployment already writes, prompt storage on or off. The two cohorts self-select, so this is directional evidence and not an experiment: a deployment that pins its hardest prompts to one model and routes only the easy ones will diff --git a/tests/test_litellm/proxy/db/test_autorouter_quality_signals.py b/tests/test_litellm/proxy/db/test_autorouter_quality_signals.py index f391d056252..dbad2433d45 100644 --- a/tests/test_litellm/proxy/db/test_autorouter_quality_signals.py +++ b/tests/test_litellm/proxy/db/test_autorouter_quality_signals.py @@ -6,6 +6,8 @@ plain Turn objects and a fixed rank table, so they exercise the actual escalatio abandonment/eligibility rules without touching Postgres or a real Router. """ +import pytest + from litellm.proxy.db.autorouter_quality_signals import ( MIN_COHORT_SESSIONS, MIN_SESSION_ID_COVERAGE, @@ -15,14 +17,24 @@ from litellm.proxy.db.autorouter_quality_signals import ( session_escalated, signals_for_cohort, ) +from litellm.router import Router # haiku < sonnet < opus, matching how rank_models_by_cost would order real deployments RANKS = {"haiku": 0, "sonnet": 1, "opus": 2} +REACHABLE_BY_KEY = {"key1": ("haiku", "sonnet", "opus")} -def _turn(session_id="s1", model="sonnet", started_at=0.0, client_disconnected=False, has_client_session_id=True): +def _turn( + session_id="s1", + model="sonnet", + started_at=0.0, + client_disconnected=False, + has_client_session_id=True, + api_key="key1", +): return Turn( session_id=session_id, + api_key=api_key, model=model, started_at=started_at, client_disconnected=client_disconnected, @@ -31,6 +43,20 @@ def _turn(session_id="s1", model="sonnet", started_at=0.0, client_disconnected=F ) +class TestRankModelsByCost: + def test_equal_cost_models_receive_equal_rank(self, monkeypatch: pytest.MonkeyPatch): + import litellm.proxy.db.autorouter_quality_signals as module + + def _fake_priced(router, candidate): + cost = {"cheap-a": 1.0, "cheap-b": 1.0, "dear": 2.0}[candidate.model] + return (cost, candidate) + + monkeypatch.setattr("litellm.router_strategy.savings_baseline._priced", _fake_priced) + ranks = module.rank_models_by_cost(Router(model_list=[]), ["cheap-a", "cheap-b", "dear"]) + assert ranks["cheap-a"] == ranks["cheap-b"], "same-cost models must share the same rank" + assert ranks["dear"] > ranks["cheap-a"], "costlier model must rank higher" + + class TestSessionEscalated: def test_upward_move_is_escalation(self): turns = [_turn(model="sonnet", started_at=1), _turn(model="opus", started_at=2)] @@ -110,7 +136,7 @@ class TestSignalsForCohort: _turn(session_id="s1", model="opus", started_at=2), _turn(session_id="s2", model="opus", started_at=1), ] - result = signals_for_cohort(turns, RANKS, reachable=["haiku", "sonnet", "opus"]) + result = signals_for_cohort(turns, RANKS, REACHABLE_BY_KEY) assert result.sessions == 1 assert result.escalation_rate_pct == 100.0 @@ -121,7 +147,7 @@ class TestSignalsForCohort: # s2 is on the ceiling model and ineligible; its disconnect must not be counted. _turn(session_id="s2", model="opus", started_at=1, client_disconnected=True), ] - result = signals_for_cohort(turns, RANKS, reachable=["haiku", "sonnet", "opus"]) + result = signals_for_cohort(turns, RANKS, REACHABLE_BY_KEY) assert result.sessions == 1 assert result.abandonment_rate_pct == 50.0 @@ -129,7 +155,7 @@ class TestSignalsForCohort: # Every session already on the ceiling model: zero would misleadingly claim # "no miss detected" when in fact nothing could be measured. turns = [_turn(session_id="s1", model="opus")] - result = signals_for_cohort(turns, RANKS, reachable=["haiku", "sonnet", "opus"]) + result = signals_for_cohort(turns, RANKS, REACHABLE_BY_KEY) assert result.sessions == 0 assert result.escalation_rate_pct is None assert result.abandonment_rate_pct is None @@ -141,10 +167,37 @@ class TestSignalsForCohort: _turn(session_id="s1", model="sonnet", started_at=1), _turn(session_id="s1", model="haiku", started_at=2), ] - result = signals_for_cohort(turns, RANKS, reachable=["haiku", "sonnet", "opus"]) + result = signals_for_cohort(turns, RANKS, REACHABLE_BY_KEY) assert result.sessions == 1 assert result.escalation_rate_pct == 0.0 + def test_same_session_id_from_two_keys_is_never_spliced_into_one_escalation(self): + # Two different callers' keys reuse the same session id. Splicing them into one + # session would read key1's haiku turn followed by key2's opus turn as one caller + # escalating -- an escalation neither caller ever made. Kept separate, each is its + # own single-turn session and neither can escalate within itself. + turns = [ + _turn(session_id="shared", api_key="key1", model="haiku", started_at=1), + _turn(session_id="shared", api_key="key2", model="opus", started_at=2), + ] + reachable = {"key1": ("haiku", "sonnet", "opus"), "key2": ("haiku", "sonnet", "opus")} + result = signals_for_cohort(turns, RANKS, reachable) + assert result.escalation_rate_pct == 0.0 + + def test_reachability_is_scoped_to_each_session_own_key(self): + # key1 can only ever reach haiku; key2 can reach up to opus. A pooled reachable set + # would let key1's session borrow key2's ceiling and count as eligible when it never + # had anywhere to escalate to. + turns = [ + _turn(session_id="s1", api_key="key1", model="haiku", started_at=1), + _turn(session_id="s2", api_key="key2", model="haiku", started_at=1), + _turn(session_id="s2", api_key="key2", model="opus", started_at=2), + ] + reachable = {"key1": ("haiku",), "key2": ("haiku", "sonnet", "opus")} + result = signals_for_cohort(turns, RANKS, reachable) + assert result.sessions == 1 + assert result.escalation_rate_pct == 100.0 + class TestBaselineUnavailableReason: def test_low_session_id_coverage_reported_before_size(self): diff --git a/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py index ef32d0f7f4a..ce801030aeb 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py @@ -616,6 +616,40 @@ class TestAutoRouterQualitySignals: response = await self._call(rows, monkeypatch) assert response.totals.baseline is None + @pytest.mark.asyncio + async def test_two_keys_reusing_the_same_session_id_are_not_merged(self, monkeypatch: pytest.MonkeyPatch): + # key-1 and key-2 both happen to send session_id="shared". key-1's own two turns + # under it escalate; key-2's own two turns under it do not (key-2's access to the + # pricey model is established by a separate "other" session so it is still eligible + # to escalate). Merging the two keys' rows by session_id alone would splice all four + # turns into one session and read key-1's escalation and key-2's non-escalation as a + # single, order-dependent sequence instead of two independent sessions. + rows = [ + self._row("shared", self.CHEAP, 1.0, api_key="key-1"), + self._row("shared", self.PRICEY, 2.0, api_key="key-1"), + self._row("shared", self.CHEAP, 1.0, api_key="key-2"), + self._row("shared", self.CHEAP, 2.0, api_key="key-2"), + self._row("other", self.PRICEY, 1.0, api_key="key-2"), + ] + response = await self._call(rows, monkeypatch) + assert response.totals.routed.sessions == 2 + assert response.totals.routed.escalation_rate_pct == 50.0 + + @pytest.mark.asyncio + async def test_reachability_is_not_pooled_across_keys_in_the_same_cohort(self, monkeypatch: pytest.MonkeyPatch): + # key-1 only ever calls the cheap model; key-2 reaches both. Pooling reachability + # across the cohort would credit key-1's session with key-2's ceiling and count it + # as eligible to escalate when key-1 never had a pricier model to reach for. + rows = [ + self._row("k1-s1", self.CHEAP, 1.0, api_key="key-1"), + self._row("k1-s1", self.CHEAP, 2.0, api_key="key-1"), + self._row("k2-s1", self.CHEAP, 1.0, api_key="key-2"), + self._row("k2-s1", self.PRICEY, 2.0, api_key="key-2"), + ] + response = await self._call(rows, monkeypatch) + assert response.totals.routed.sessions == 1 + assert response.totals.routed.escalation_rate_pct == 100.0 + @pytest.mark.asyncio async def test_groups_are_reported_per_router(self, monkeypatch: pytest.MonkeyPatch): rows = [