From 1bba6094c5c40e0dc2191b0bc62a84f2a370948f Mon Sep 17 00:00:00 2001 From: Abhimanyu Kapur <38531241+akapur99@users.noreply.github.com> Date: Fri, 7 Aug 2026 09:52:08 -0700 Subject: [PATCH] feat(adaptive_router): add efficiency posterior and latency/throughput reward MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Ports the efficiency-reward piece of the Elastic Unified Router (EURo) design into the adaptive router's bandit, giving it a second Beta posterior keyed on latency and completion-token throughput, independent of quality signals. ## What changed - BanditCell now tracks (alpha_eff, beta_eff) alongside (alpha, beta), both fed independently. Efficiency defaults to Beta(1, 1) prior (uninformative). - AdaptiveRouter.record_efficiency() updates the efficiency posterior from observed latency and token count; independent of record_turn()'s quality signals. - Efficiency reward = gamma * exp(-latency/target) + (1-gamma) * min(tokens/sec/max, 1), clamped to [0, 1]. On failure (non-2xx), reward = 0. - Selection now blends three terms: quality_weight * TS(quality) + efficiency_weight * TS(efficiency) + cost_weight * cost. Efficiency weight defaults to 0.0 (opt-in). - DB: schema.prisma + migration add alpha_eff/beta_eff columns to LiteLLM_AdaptiveRouterState; read/write plumbing in update_queue and adaptive_router. - Snapshot endpoint reports efficiency_mean alongside quality_mean. ## Tests New tests for apply_efficiency_delta, efficiency_mean property, and flush behavior with efficiency deltas. All existing bandit and router tests pass. ## Design decisions baked in (as per PHASE1_PLAN.md) 1. Efficiency weight defaults to 0.0 — no silent behavior change for existing deployments. Opt-in only. 2. target_latency is static (1.0s default, overridable in config). Auto-calibration deferred to Phase 2. 3. Failure handling reuses existing response_status detection — one source of truth. 4. Schema extended in place (not a sibling table) — (router, request_type, model) is already the right key for both quality and efficiency. Co-Authored-By: Claude --- .../migration.sql | 4 + .../litellm_proxy_extras/schema.prisma | 6 + litellm/proxy/schema.prisma | 6 + .../router_strategy/adaptive_router/README.md | 45 +++++--- .../adaptive_router/adaptive_router.py | 68 ++++++++++- .../router_strategy/adaptive_router/bandit.py | 60 +++++++--- .../router_strategy/adaptive_router/config.py | 7 ++ .../adaptive_router/efficiency.py | 29 +++++ .../router_strategy/adaptive_router/hooks.py | 48 +++++++- .../adaptive_router/update_queue.py | 14 +++ .../complexity_router/complexity_router.py | 13 ++- litellm/types/router.py | 10 +- schema.prisma | 6 + .../adaptive_router/test_adaptive_router.py | 109 ++++++++++++++++++ .../adaptive_router/test_bandit.py | 22 ++++ .../adaptive_router/test_state_endpoint.py | 5 +- .../adaptive_router/test_update_queue.py | 25 ++++ 17 files changed, 441 insertions(+), 36 deletions(-) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260807094919_add_efficiency_posterior_to_adaptive_router/migration.sql create mode 100644 litellm/router_strategy/adaptive_router/efficiency.py diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260807094919_add_efficiency_posterior_to_adaptive_router/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260807094919_add_efficiency_posterior_to_adaptive_router/migration.sql new file mode 100644 index 00000000000..702b2e698a6 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260807094919_add_efficiency_posterior_to_adaptive_router/migration.sql @@ -0,0 +1,4 @@ +-- AlterTable +ALTER TABLE "LiteLLM_AdaptiveRouterState" ADD COLUMN "alpha_eff" DOUBLE PRECISION NOT NULL DEFAULT 1.0, +ADD COLUMN "beta_eff" DOUBLE PRECISION NOT NULL DEFAULT 1.0; + diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 9c871b65f40..d015b7f7804 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -1377,6 +1377,12 @@ model LiteLLM_AdaptiveRouterState { model_name String alpha Float beta Float + // Efficiency posterior (latency/throughput reward) -- independent Beta on the same + // cell, fed by AdaptiveRouter.record_efficiency. Defaults to the uninformative + // Beta(1, 1) prior used in-memory (see bandit.py) so a row created before this + // column existed reads back with no efficiency signal yet, not a false one. + alpha_eff Float @default(1.0) + beta_eff Float @default(1.0) total_samples Int @default(0) last_updated_at DateTime @default(now()) @updatedAt diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 9c871b65f40..d015b7f7804 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -1377,6 +1377,12 @@ model LiteLLM_AdaptiveRouterState { model_name String alpha Float beta Float + // Efficiency posterior (latency/throughput reward) -- independent Beta on the same + // cell, fed by AdaptiveRouter.record_efficiency. Defaults to the uninformative + // Beta(1, 1) prior used in-memory (see bandit.py) so a row created before this + // column existed reads back with no efficiency signal yet, not a false one. + alpha_eff Float @default(1.0) + beta_eff Float @default(1.0) total_samples Int @default(0) last_updated_at DateTime @default(now()) @updatedAt diff --git a/litellm/router_strategy/adaptive_router/README.md b/litellm/router_strategy/adaptive_router/README.md index 09420a8dd9d..e1f427db2b1 100644 --- a/litellm/router_strategy/adaptive_router/README.md +++ b/litellm/router_strategy/adaptive_router/README.md @@ -1,15 +1,17 @@ -# Adaptive Router (v0) +# Adaptive Router (v0.1) A request-type-aware routing strategy. For each incoming request, classify the prompt into one of seven `RequestType` buckets (code generation, writing, analytical reasoning, …), then Thompson-sample a Beta(α, β) bandit posterior -per `(request_type, model)` cell to pick the best model. Quality estimates are -combined with a normalized cost score via a weighted linear sum. +per `(request_type, model)` cell to pick the best model. Quality estimates, +efficiency estimates (latency + throughput), and normalized cost are combined +via a weighted linear sum. A post-call hook reads the response and runs lightweight regex + tool-call detectors (see `signals.py`) to award per-turn credit/blame to the model that -served the turn. Updates are batched in-memory and flushed to Postgres every -~10s by a background task in `proxy_server.py`. +served the turn. Independently, it extracts latency and token count to update +an efficiency posterior (see `efficiency.py`). Updates are batched in-memory +and flushed to Postgres every ~10s by a background task in `proxy_server.py`. ## Config example @@ -42,6 +44,7 @@ model_list: weights: quality: 0.7 cost: 0.3 + efficiency: 0.0 # Optional: set to >0 to factor latency/throughput into selection ``` Callers may pass header `x-litellm-min-quality-tier: 3` (or metadata key @@ -54,27 +57,43 @@ Callers may pass header `x-litellm-min-quality-tier: 3` (or metadata key and total mass = `COLD_START_MASS` (10). About ten real observations move it meaningfully. - **Per-request decision.** Sample once per eligible model, score with - `quality_weight·sample + cost_weight·normalized_cost`, pick the argmax. - Routing is stateless per-turn — no sticky lookup. Each call resamples. + `quality_weight·TS(quality) + efficiency_weight·TS(efficiency) + cost_weight·normalized_cost`, + pick the argmax. Efficiency is `gamma·exp(-latency/target) + (1-gamma)·min(tokens/sec/max, 1)`; + see `efficiency.py`. Routing is stateless per-turn — no sticky lookup. Each call resamples. - **Previous-response attribution.** Post-call, feedback from the current user message is attributed to the model that produced the previous response, while response signals are attributed to the current model. Contexts expire after 24 hours and the in-memory cache is capped at 1,024 sessions. Conversation identity is the client-supplied `litellm_session_id` if present, otherwise a sha256 over caller identity (api key hash, team, user, end-user) + the first message. -- **Per-turn updates.** `satisfaction → +α`. `misalignment, stagnation, +- **Per-turn quality updates.** `satisfaction → +α`. `misalignment, stagnation, disengagement, failure → +β` (each). `loop → +0.5β`. `exhaustion → 0` (uptime, not quality). Skipped if conversation has fewer than `SIGNAL_GATE_MIN_MESSAGES` messages. +- **Per-response efficiency updates.** Every completion updates the efficiency + posterior independently: `reward = gamma·exp(-latency/target_latency) + (1-gamma)·throughput`, + then `+α_eff = reward`, `+β_eff = 1 - reward`. No cold-start prior; starts at Beta(1, 1). + On failure (non-2xx), `reward = 0` (same failure detection as quality signals). - **Persistence.** Bandit cells: aggregated deltas, eventually consistent. + Both quality and efficiency posteriors live in the same row in `LiteLLM_AdaptiveRouterState`. Session rows: last-write-wins snapshots. -## Known v0 limitations +## Efficiency weighting (v0.1+) -- **Latency is not in the score.** Quality + cost only. A pathologically slow - model can still be picked. -- **Hard sample cap at 200.** Once `α + β > 200`, deltas are silently dropped. - No rescaling — drift is a v1 concern. +The `efficiency` weight is **optional and defaults to 0.0** (off by default). +When set to a value > 0, it reallocates from the `quality` + `cost` budget: +the three weights must still sum to 1.0. Operators should start with a small +value (e.g., `efficiency: 0.1`) and tune based on their workload's latency +sensitivity. No auto-calibration yet; `target_latency` is a static 1.0s by default +and can be overridden in config. + +## Known v0.1 limitations + +- **Hard sample cap at 200.** Once `α + β > 200` (quality or efficiency), deltas + are silently dropped. No rescaling — drift is a v1 concern. This cap is shared + between quality and efficiency on the same cell. +- **No auto-calibration.** `target_latency` is static (1.0s default). Phase 2 + will add a warm-up window to auto-detect per model group. - **Signals are regex + tool-call only.** No LLM-judge, no embedding similarity, no exemplar storage. Signals are best-effort and biased toward English. - **One AdaptiveRouter per `Router`.** Multiple `adaptive_router/*` deployments diff --git a/litellm/router_strategy/adaptive_router/adaptive_router.py b/litellm/router_strategy/adaptive_router/adaptive_router.py index 1a33ea23bd4..d9606c761ee 100644 --- a/litellm/router_strategy/adaptive_router/adaptive_router.py +++ b/litellm/router_strategy/adaptive_router/adaptive_router.py @@ -26,10 +26,12 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( from litellm.router_strategy.adaptive_router.bandit import ( BanditCell, apply_delta, + apply_efficiency_delta, initial_cell, pick_best, ) from litellm.router_strategy.adaptive_router.classifier import classify_prompt +from litellm.router_strategy.adaptive_router.efficiency import composite_efficiency_reward from litellm.router_strategy.adaptive_router.config import ( ADAPTIVE_ROUTER_CHOSEN_MODEL_KEY, MIN_QUALITY_TIER_HEADER, @@ -139,7 +141,12 @@ class AdaptiveRouter: continue if row.model_name not in self.config.available_models: continue - self._cells[(rt, row.model_name)] = BanditCell(alpha=row.alpha, beta=row.beta) + self._cells[(rt, row.model_name)] = BanditCell( + alpha=row.alpha, + beta=row.beta, + alpha_eff=getattr(row, "alpha_eff", 1.0), + beta_eff=getattr(row, "beta_eff", 1.0), + ) loaded += 1 verbose_router_logger.info( "AdaptiveRouter[%s]: loaded %d cells from DB", @@ -225,6 +232,7 @@ class AdaptiveRouter: costs, quality_weight=self.config.weights.quality, cost_weight=self.config.weights.cost, + efficiency_weight=self.config.weights.efficiency, ) async def get_state_snapshot(self) -> dict[str, Any]: @@ -232,18 +240,22 @@ class AdaptiveRouter: cells: Final = [] for (rt, model), cell in sorted(self._cells.items(), key=lambda kv: (kv[0][0].value, kv[0][1])): total = cell.alpha + cell.beta + total_eff = cell.alpha_eff + cell.beta_eff cells.append( { "request_type": rt.value, "model": model, "alpha": cell.alpha, "beta": cell.beta, + "alpha_eff": cell.alpha_eff, + "beta_eff": cell.beta_eff, # Net observations that have moved the posterior, excluding # the cold-start prior mass. `alpha + beta` would show the # initial COLD_START_MASS (e.g. 10) before any real traffic # arrives, which confuses operators reading the endpoint. "samples": cell.total_samples, "quality_mean": cell.alpha / total if total > 0 else 0.0, + "efficiency_mean": cell.alpha_eff / total_eff if total_eff > 0 else 0.0, } ) queue: Final = await self.queue.queue_size() @@ -255,6 +267,7 @@ class AdaptiveRouter: "weights": { "quality": self.config.weights.quality, "cost": self.config.weights.cost, + "efficiency": self.config.weights.efficiency, }, "model_costs": dict(self.model_to_cost), "cells": cells, @@ -473,6 +486,59 @@ class AdaptiveRouter: ) return combined_delta + async def record_efficiency( + self, + model_name: str, + request_type: RequestType, + latency_seconds: float, + completion_tokens: int, + is_failure: bool = False, + ) -> float: + """ + Update the (request_type, model) cell's efficiency posterior from one observed + call. Returns the reward that was applied. + """ + reward: Final = ( + 0.0 + if is_failure + else composite_efficiency_reward( + latency_seconds=latency_seconds, + completion_tokens=completion_tokens, + ) + ) + d_alpha_eff: Final = reward + d_beta_eff: Final = 1.0 - reward + async with self._lock: + cell_key: Final = (request_type, model_name) + if cell_key not in self._cells: + self._cells[cell_key] = initial_cell( + self.model_to_prefs.get(model_name, _default_prefs()), request_type + ) + before: Final = self._cells[cell_key] + after: Final = apply_efficiency_delta(before, d_alpha_eff, d_beta_eff) + self._cells[cell_key] = after + if after is before: + return reward + await self.queue.add_state_delta( + self.router_name, + request_type.value, + model_name, + delta_alpha=0.0, + delta_beta=0.0, + delta_alpha_eff=d_alpha_eff, + delta_beta_eff=d_beta_eff, + ) + verbose_router_logger.debug( + "AdaptiveRouter[%s]: efficiency model=%s request_type=%s latency=%.3fs tokens=%d reward=%.3f", + self.router_name, + model_name, + request_type.value, + latency_seconds, + completion_tokens, + reward, + ) + return reward + @staticmethod def _persistable_session_snapshot(state: SessionState) -> dict[str, Any]: snapshot: Final = asdict(state) diff --git a/litellm/router_strategy/adaptive_router/bandit.py b/litellm/router_strategy/adaptive_router/bandit.py index c2eb232bd99..6f12124f0da 100644 --- a/litellm/router_strategy/adaptive_router/bandit.py +++ b/litellm/router_strategy/adaptive_router/bandit.py @@ -1,13 +1,5 @@ """ Thompson sampling and prior initialization for the adaptive router bandit. - -Each (router, request_type, model) cell is a Beta(alpha, beta) posterior. -- alpha = pseudo-successes -- beta = pseudo-failures -- mean = alpha / (alpha + beta) -- total samples = alpha + beta - COLD_START_MASS (informative prior, not data) - -Hot path: thompson_sample() — pure function, no I/O. """ import random @@ -18,7 +10,10 @@ from litellm.router_strategy.adaptive_router.config import ( BASE_TIER_WEIGHT, COLD_START_MASS, DEFAULT_COST_WEIGHT, + DEFAULT_EFFICIENCY_WEIGHT, DEFAULT_QUALITY_WEIGHT, + EFFICIENCY_PRIOR_ALPHA, + EFFICIENCY_PRIOR_BETA, SAMPLE_CAP, STRENGTH_BONUS, ) @@ -31,12 +26,19 @@ class BanditCell: alpha: float beta: float + alpha_eff: float = EFFICIENCY_PRIOR_ALPHA + beta_eff: float = EFFICIENCY_PRIOR_BETA @property def mean(self) -> float: total: Final = self.alpha + self.beta return self.alpha / total if total > 0 else 0.5 + @property + def efficiency_mean(self) -> float: + total: Final = self.alpha_eff + self.beta_eff + return self.alpha_eff / total if total > 0 else 0.5 + @property def total_samples(self) -> int: return max(0, int(self.alpha + self.beta - COLD_START_MASS)) @@ -63,7 +65,7 @@ def initial_cell(prefs: AdaptiveRouterPreferences, request_type: RequestType) -> def apply_delta(cell: BanditCell, delta_alpha: float, delta_beta: float) -> BanditCell: """ - Apply a learning update to a cell, enforcing the sample cap. + Apply a learning update to the quality posterior, enforcing the sample cap. SAMPLE_CAP is a HARD cap on (alpha + beta). When the cap would be exceeded, we drop the update. (D5: hard cap, no rescaling — keep v0 simple.) @@ -72,15 +74,30 @@ def apply_delta(cell: BanditCell, delta_alpha: float, delta_beta: float) -> Band new_beta: Final = cell.beta + delta_beta if new_alpha + new_beta > SAMPLE_CAP: return cell - return BanditCell(alpha=new_alpha, beta=new_beta) + return BanditCell(alpha=new_alpha, beta=new_beta, alpha_eff=cell.alpha_eff, beta_eff=cell.beta_eff) + + +def apply_efficiency_delta(cell: BanditCell, delta_alpha_eff: float, delta_beta_eff: float) -> BanditCell: + """Apply a learning update to the efficiency posterior, enforcing the same sample cap.""" + new_alpha_eff: Final = cell.alpha_eff + delta_alpha_eff + new_beta_eff: Final = cell.beta_eff + delta_beta_eff + if new_alpha_eff + new_beta_eff > SAMPLE_CAP: + return cell + return BanditCell(alpha=cell.alpha, beta=cell.beta, alpha_eff=new_alpha_eff, beta_eff=new_beta_eff) def thompson_sample(cell: BanditCell, rng: random.Random | None = None) -> float: - """Draw a sample from Beta(alpha, beta). Returns a quality estimate in [0, 1].""" + """Draw a sample from the quality posterior Beta(alpha, beta). Returns an estimate in [0, 1].""" r: Final = rng if rng is not None else random return r.betavariate(cell.alpha, cell.beta) +def thompson_sample_efficiency(cell: BanditCell, rng: random.Random | None = None) -> float: + """Draw a sample from the efficiency posterior Beta(alpha_eff, beta_eff).""" + r: Final = rng if rng is not None else random + return r.betavariate(cell.alpha_eff, cell.beta_eff) + + def normalized_cost(model_cost: float, all_costs: list[float]) -> float: """ Map a raw $/1k-token cost into [0, 1] where 0 = most expensive, 1 = cheapest. @@ -100,13 +117,20 @@ def score( all_costs: list[float], quality_weight: float = DEFAULT_QUALITY_WEIGHT, cost_weight: float = DEFAULT_COST_WEIGHT, + efficiency_sample: float | None = None, + efficiency_weight: float = DEFAULT_EFFICIENCY_WEIGHT, ) -> float: """ - Multi-objective score. V0 is a weighted linear sum of (quality, normalized_cost). - Higher is better. Both inputs are in [0, 1]. + Multi-objective score. V0 was a weighted linear sum of (quality, normalized_cost); + efficiency is an additive third term so a caller that never passes efficiency_sample + (efficiency_weight defaults to 0.0) gets byte-identical scores to before this was added. + Higher is better. All inputs are in [0, 1]. """ cost_score: Final = normalized_cost(model_cost, all_costs) - return quality_weight * quality_sample + cost_weight * cost_score + total: Final = quality_weight * quality_sample + cost_weight * cost_score + if efficiency_sample is None or efficiency_weight == 0.0: + return total + return total + efficiency_weight * efficiency_sample def pick_best( @@ -114,6 +138,7 @@ def pick_best( model_costs: dict[str, float], quality_weight: float = DEFAULT_QUALITY_WEIGHT, cost_weight: float = DEFAULT_COST_WEIGHT, + efficiency_weight: float = DEFAULT_EFFICIENCY_WEIGHT, rng: random.Random | None = None, ) -> str: """ @@ -121,6 +146,10 @@ def pick_best( cells: {model_name: BanditCell} model_costs: {model_name: $/1k tokens} + + Efficiency is only sampled when efficiency_weight != 0.0, so a caller that never + passes it (the default) does one fewer betavariate draw per candidate -- same hot + path cost as before this was added, not just the same score. """ if not cells: raise ValueError("pick_best called with no models") @@ -129,7 +158,8 @@ def pick_best( best_score = float("-inf") for model, cell in cells.items(): q = thompson_sample(cell, rng=rng) - s = score(q, model_costs[model], all_costs, quality_weight, cost_weight) + eff = thompson_sample_efficiency(cell, rng=rng) if efficiency_weight != 0.0 else None + s = score(q, model_costs[model], all_costs, quality_weight, cost_weight, eff, efficiency_weight) if s > best_score: best_score = s best_model = model diff --git a/litellm/router_strategy/adaptive_router/config.py b/litellm/router_strategy/adaptive_router/config.py index af195b42f84..2ba0f1d41eb 100644 --- a/litellm/router_strategy/adaptive_router/config.py +++ b/litellm/router_strategy/adaptive_router/config.py @@ -13,6 +13,13 @@ from litellm.types.router import RequestType # re-export for convenience # noq DEFAULT_QUALITY_WEIGHT: Final[float] = 0.7 # UNVALIDATED — calibrated against [0] sessions DEFAULT_COST_WEIGHT: Final[float] = 0.3 # UNVALIDATED — calibrated against [0] sessions +DEFAULT_EFFICIENCY_WEIGHT: Final[float] = 0.0 +EFFICIENCY_GAMMA: Final[float] = 0.5 +DEFAULT_TARGET_LATENCY_SECONDS: Final[float] = 1.0 +EFFICIENCY_THROUGHPUT_MAX: Final[float] = 100.0 +EFFICIENCY_PRIOR_ALPHA: Final[float] = 1.0 +EFFICIENCY_PRIOR_BETA: Final[float] = 1.0 + # D4 — Cold-start prior: (alpha + beta) total mass = COLD_START_MASS # Mean of Beta = base_tier_weight + (strength_bonus if declared) BASE_TIER_WEIGHT: Final[dict[int, float]] = {1: 0.3, 2: 0.5, 3: 0.7} # UNVALIDATED diff --git a/litellm/router_strategy/adaptive_router/efficiency.py b/litellm/router_strategy/adaptive_router/efficiency.py new file mode 100644 index 00000000000..4613a0cb0e4 --- /dev/null +++ b/litellm/router_strategy/adaptive_router/efficiency.py @@ -0,0 +1,29 @@ +""" +Efficiency reward for the adaptive router's (request_type, model) bandit cells. +""" + +from __future__ import annotations + +import math +from typing import Final + +from litellm.router_strategy.adaptive_router.config import ( + DEFAULT_TARGET_LATENCY_SECONDS, + EFFICIENCY_GAMMA, + EFFICIENCY_THROUGHPUT_MAX, +) + + +def composite_efficiency_reward( + latency_seconds: float, + completion_tokens: int, + gamma: float = EFFICIENCY_GAMMA, + target_latency_seconds: float = DEFAULT_TARGET_LATENCY_SECONDS, + throughput_max: float = EFFICIENCY_THROUGHPUT_MAX, +) -> float: + """Blend a latency-decay score and a throughput score into a single reward in [0, 1].""" + safe_latency: Final = max(latency_seconds, 1e-6) + latency_score: Final = math.exp(-safe_latency / max(target_latency_seconds, 1e-6)) + throughput_score: Final = min(max(completion_tokens, 0) / safe_latency / max(throughput_max, 1e-6), 1.0) + reward: Final = gamma * latency_score + (1.0 - gamma) * throughput_score + return min(max(reward, 0.0), 1.0) diff --git a/litellm/router_strategy/adaptive_router/hooks.py b/litellm/router_strategy/adaptive_router/hooks.py index b59ce6e3621..69d8bfe0d2c 100644 --- a/litellm/router_strategy/adaptive_router/hooks.py +++ b/litellm/router_strategy/adaptive_router/hooks.py @@ -13,6 +13,7 @@ from __future__ import annotations import hashlib import json +from datetime import timedelta from typing import Any, Final from litellm._logging import verbose_router_logger @@ -25,6 +26,7 @@ from litellm.router_strategy.adaptive_router.config import ( SIGNAL_GATE_MIN_MESSAGES, ) from litellm.router_strategy.adaptive_router.signals import Turn +from litellm.types.utils import ModelResponse # Identity fields hashed into a derived session key so the same conversation # from the same caller produces a stable key, while different keys/teams/users @@ -197,20 +199,50 @@ class AdaptiveRouterPostCallHook(CustomLogger): return {ADAPTIVE_ROUTER_RESPONSE_HEADER: chosen} async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): - await self._record(kwargs, response_obj, response_status=200) + await self._record(kwargs, response_obj, response_status=200, start_time=start_time, end_time=end_time) async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): status = kwargs.get("response_status") if status is None: exc: Final = kwargs.get("exception") status = getattr(exc, "status_code", 500) if exc is not None else 500 - await self._record(kwargs, response_obj, response_status=int(status)) + await self._record( + kwargs, response_obj, response_status=int(status), start_time=start_time, end_time=end_time + ) + + @staticmethod + def _extract_latency_seconds(start_time: Any, end_time: Any) -> float | None: + if start_time is None or end_time is None: + return None + try: + elapsed = end_time - start_time + except TypeError: + return None + if isinstance(elapsed, timedelta): + elapsed = elapsed.total_seconds() + try: + elapsed = float(elapsed) + except (TypeError, ValueError): + return None + return elapsed if elapsed > 0 else None + + @staticmethod + def _extract_completion_tokens(response_obj: Any) -> int | None: + if not isinstance(response_obj, ModelResponse): + return None + usage: Final = getattr(response_obj, "usage", None) + if usage is None: + return None + tokens = getattr(usage, "completion_tokens", None) + return int(tokens) if isinstance(tokens, (int, float)) else None async def _record( self, kwargs: dict[str, Any], response_obj: Any, response_status: int, + start_time: Any = None, + end_time: Any = None, ) -> None: try: messages: Final = kwargs.get("messages") or [] @@ -247,5 +279,17 @@ class AdaptiveRouterPostCallHook(CustomLogger): request_type=request_type, turn=turn, ) + + latency_seconds: Final = self._extract_latency_seconds(start_time, end_time) + if latency_seconds is not None: + is_failure: Final = not (200 <= response_status < 300) + completion_tokens: Final = 0 if is_failure else (self._extract_completion_tokens(response_obj) or 0) + await self.adaptive_router.record_efficiency( + model_name=current_model, + request_type=request_type, + latency_seconds=latency_seconds, + completion_tokens=completion_tokens, + is_failure=is_failure, + ) except Exception as e: verbose_router_logger.exception("AdaptiveRouterPostCallHook: failed to record turn: %s", e) diff --git a/litellm/router_strategy/adaptive_router/update_queue.py b/litellm/router_strategy/adaptive_router/update_queue.py index 1b9fce284ac..84b46cab34e 100644 --- a/litellm/router_strategy/adaptive_router/update_queue.py +++ b/litellm/router_strategy/adaptive_router/update_queue.py @@ -26,6 +26,10 @@ from litellm.repositories.table_repositories import ( AdaptiveRouterSessionRepository, AdaptiveRouterStateRepository, ) +from litellm.router_strategy.adaptive_router.config import ( + EFFICIENCY_PRIOR_ALPHA, + EFFICIENCY_PRIOR_BETA, +) StateKey = tuple[str, str, str] # (router_name, request_type, model_name) SessionKey = tuple[str, str, str] # (session_id, router_name, model_name) @@ -53,6 +57,8 @@ class AdaptiveRouterUpdateQueue: model_name: str, delta_alpha: float, delta_beta: float, + delta_alpha_eff: float = 0.0, + delta_beta_eff: float = 0.0, ) -> None: """Aggregate a bandit-cell delta. Multiple deltas to the same cell sum.""" key: Final[StateKey] = (router_name, request_type, model_name) @@ -62,11 +68,15 @@ class AdaptiveRouterUpdateQueue: self._state_agg[key] = { "delta_alpha": delta_alpha, "delta_beta": delta_beta, + "delta_alpha_eff": delta_alpha_eff, + "delta_beta_eff": delta_beta_eff, "samples_added": 1, } else: current["delta_alpha"] += delta_alpha current["delta_beta"] += delta_beta + current["delta_alpha_eff"] += delta_alpha_eff + current["delta_beta_eff"] += delta_beta_eff current["samples_added"] += 1 self._max_state_size_seen = max(self._max_state_size_seen, len(self._state_agg)) @@ -129,11 +139,15 @@ class AdaptiveRouterUpdateQueue: "model_name": model, "alpha": payload["delta_alpha"], "beta": payload["delta_beta"], + "alpha_eff": EFFICIENCY_PRIOR_ALPHA + payload["delta_alpha_eff"], + "beta_eff": EFFICIENCY_PRIOR_BETA + payload["delta_beta_eff"], "total_samples": int(payload["samples_added"]), }, "update": { "alpha": {"increment": payload["delta_alpha"]}, "beta": {"increment": payload["delta_beta"]}, + "alpha_eff": {"increment": payload["delta_alpha_eff"]}, + "beta_eff": {"increment": payload["delta_beta_eff"]}, "total_samples": {"increment": int(payload["samples_added"])}, }, }, diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index 32d252f3f68..5cdc9e8384b 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -1294,6 +1294,7 @@ class ComplexityRouter(CustomLogger): from litellm.router_strategy.adaptive_router.bandit import ( normalized_cost, thompson_sample, + thompson_sample_efficiency, ) from litellm.router_strategy.adaptive_router.classifier import classify_prompt @@ -1320,6 +1321,7 @@ class ComplexityRouter(CustomLogger): "eligible_mode": "classified_tier", "quality_weight": self.config.adaptive_weights.quality, "cost_weight": self.config.adaptive_weights.cost, + "efficiency_weight": self.config.adaptive_weights.efficiency, "tier_distance_penalty": self.config.tier_distance_penalty, "chosen_model": chosen_model, "candidates": [ @@ -1341,6 +1343,7 @@ class ComplexityRouter(CustomLogger): all_costs: Final = [adaptive.model_to_cost.get(m, 0.0) for m in candidates] quality_weight: Final = self.config.adaptive_weights.quality cost_weight: Final = self.config.adaptive_weights.cost + efficiency_weight: Final = self.config.adaptive_weights.efficiency penalty_weight: Final = self.config.tier_distance_penalty best_model: str | None = None @@ -1350,6 +1353,7 @@ class ComplexityRouter(CustomLogger): cell = adaptive._cells[(request_type, model)] quality_sample = thompson_sample(cell) cost_score = normalized_cost(adaptive.model_to_cost.get(model, 0.0), all_costs) + efficiency_sample = thompson_sample_efficiency(cell) if efficiency_weight != 0.0 else 0.0 if self.config.adaptive_eligible == "classified_tier": distance = 0 else: @@ -1357,12 +1361,18 @@ class ComplexityRouter(CustomLogger): distance = min( abs(TIER_SEVERITY_ORDER.index(model_tier) - classified_idx) for model_tier in model_tiers ) - score = quality_weight * quality_sample + cost_weight * cost_score - penalty_weight * distance + score = ( + quality_weight * quality_sample + + cost_weight * cost_score + + efficiency_weight * efficiency_sample + - penalty_weight * distance + ) candidate_scores.append( { "model": model, "quality_sample": quality_sample, "cost_score": cost_score, + "efficiency_sample": efficiency_sample, "tier_distance": distance, "score": score, } @@ -1382,6 +1392,7 @@ class ComplexityRouter(CustomLogger): "eligible_mode": self.config.adaptive_eligible, "quality_weight": quality_weight, "cost_weight": cost_weight, + "efficiency_weight": efficiency_weight, "tier_distance_penalty": penalty_weight, "chosen_model": best_model, "candidates": candidate_scores, diff --git a/litellm/types/router.py b/litellm/types/router.py index e166d844735..77717092a15 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -889,13 +889,17 @@ class RequestType(str, enum.Enum): class AdaptiveRouterWeights(BaseModel): quality: float = Field(default=0.7, ge=0.0, le=1.0) cost: float = Field(default=0.3, ge=0.0, le=1.0) + efficiency: float = Field(default=0.0, ge=0.0, le=1.0, validate_default=True) - @field_validator("cost") + @field_validator("efficiency") @classmethod def _weights_sum_to_one(cls, v, info): q: Final = info.data.get("quality", 0.7) - if abs(q + v - 1.0) > 0.001: - raise ValueError(f"weights must sum to 1.0, got quality={q} + cost={v} = {q + v}") + c: Final = info.data.get("cost", 0.3) + if abs(q + c + v - 1.0) > 0.001: + raise ValueError( + f"weights must sum to 1.0, got quality={q} + cost={c} + efficiency={v} = {q + c + v}" + ) return v diff --git a/schema.prisma b/schema.prisma index 9c871b65f40..d015b7f7804 100644 --- a/schema.prisma +++ b/schema.prisma @@ -1377,6 +1377,12 @@ model LiteLLM_AdaptiveRouterState { model_name String alpha Float beta Float + // Efficiency posterior (latency/throughput reward) -- independent Beta on the same + // cell, fed by AdaptiveRouter.record_efficiency. Defaults to the uninformative + // Beta(1, 1) prior used in-memory (see bandit.py) so a row created before this + // column existed reads back with no efficiency signal yet, not a false one. + alpha_eff Float @default(1.0) + beta_eff Float @default(1.0) total_samples Int @default(0) last_updated_at DateTime @default(now()) @updatedAt diff --git a/tests/test_litellm/router_strategy/adaptive_router/test_adaptive_router.py b/tests/test_litellm/router_strategy/adaptive_router/test_adaptive_router.py index cbf5635a5ae..a399f909cba 100644 --- a/tests/test_litellm/router_strategy/adaptive_router/test_adaptive_router.py +++ b/tests/test_litellm/router_strategy/adaptive_router/test_adaptive_router.py @@ -7,6 +7,8 @@ from litellm.router_strategy.adaptive_router import adaptive_router as ar_module import pytest from litellm.router_strategy.adaptive_router.adaptive_router import AdaptiveRouter +from litellm.router_strategy.adaptive_router.bandit import BanditCell +from litellm.router_strategy.adaptive_router.config import SAMPLE_CAP from litellm.router_strategy.adaptive_router.signals import Turn from litellm.types.router import ( AdaptiveRouterConfig, @@ -268,6 +270,73 @@ async def test_record_turn_bounds_feedback_contexts_and_evicts_least_recent_sess assert "overflow" in r._feedback_contexts +@pytest.mark.asyncio +async def test_record_efficiency_pushes_to_queue(): + r = _make_router() + r.queue.add_state_delta = AsyncMock() + + reward = await r.record_efficiency( + model_name="fast", + request_type=RequestType.GENERAL, + latency_seconds=0.5, + completion_tokens=50, + ) + + assert 0.0 <= reward <= 1.0 + r.queue.add_state_delta.assert_awaited_once() + call_kwargs = r.queue.add_state_delta.call_args.kwargs + assert call_kwargs["delta_alpha_eff"] == reward + assert call_kwargs["delta_beta_eff"] == 1.0 - reward + assert call_kwargs["delta_alpha"] == 0.0 + assert call_kwargs["delta_beta"] == 0.0 + + +@pytest.mark.asyncio +async def test_record_efficiency_zeroes_reward_on_failure(): + r = _make_router() + r.queue.add_state_delta = AsyncMock() + + reward = await r.record_efficiency( + model_name="fast", + request_type=RequestType.GENERAL, + latency_seconds=0.1, + completion_tokens=50, + is_failure=True, + ) + + assert reward == 0.0 + call_kwargs = r.queue.add_state_delta.call_args.kwargs + assert call_kwargs["delta_alpha_eff"] == 0.0 + assert call_kwargs["delta_beta_eff"] == 1.0 + + +@pytest.mark.asyncio +async def test_record_efficiency_at_sample_cap_does_not_persist_rejected_update(): + """Regression: when the in-memory efficiency posterior is at SAMPLE_CAP, apply_efficiency_delta + rejects the update and leaves the cell unchanged. record_efficiency must not enqueue a DB write + for a delta that was never actually applied in memory -- otherwise the DB and in-memory state + diverge, and the row can persist an efficiency posterior the live cell never had.""" + r = _make_router() + cell_key = (RequestType.GENERAL, "fast") + r._cells[cell_key] = BanditCell( + alpha=r._cells[cell_key].alpha, + beta=r._cells[cell_key].beta, + alpha_eff=SAMPLE_CAP - 1.0, + beta_eff=1.0, + ) + r.queue.add_state_delta = AsyncMock() + + await r.record_efficiency( + model_name="fast", + request_type=RequestType.GENERAL, + latency_seconds=0.5, + completion_tokens=50, + ) + + r.queue.add_state_delta.assert_not_awaited() + assert r._cells[cell_key].alpha_eff == SAMPLE_CAP - 1.0 + + @pytest.mark.asyncio async def test_load_state_from_db_overrides_cold_start(): r = _make_router() @@ -288,6 +357,46 @@ async def test_load_state_from_db_overrides_cold_start(): assert (new_cell.alpha, new_cell.beta) != (cold.alpha, cold.beta) +@pytest.mark.asyncio +async def test_load_state_from_db_loads_efficiency_posterior(): + r = _make_router() + + fake_row = MagicMock() + fake_row.request_type = "general" + fake_row.model_name = "fast" + fake_row.alpha = 42.0 + fake_row.beta = 13.0 + fake_row.alpha_eff = 6.0 + fake_row.beta_eff = 4.0 + + prisma = MagicMock() + prisma.db.litellm_adaptiverouterstate.find_many = AsyncMock(return_value=[fake_row]) + await r.load_state_from_db(prisma) + + new_cell = r._cells[(RequestType.GENERAL, "fast")] + assert (new_cell.alpha_eff, new_cell.beta_eff) == (6.0, 4.0) + + +@pytest.mark.asyncio +async def test_load_state_from_db_defaults_efficiency_when_row_lacks_columns(): + """A row from before alpha_eff/beta_eff existed should load as the uninformative + Beta(1, 1) prior, not crash or silently carry over an unrelated attribute.""" + r = _make_router() + + class _LegacyRow: + request_type = "general" + model_name = "fast" + alpha = 42.0 + beta = 13.0 + + prisma = MagicMock() + prisma.db.litellm_adaptiverouterstate.find_many = AsyncMock(return_value=[_LegacyRow()]) + await r.load_state_from_db(prisma) + + new_cell = r._cells[(RequestType.GENERAL, "fast")] + assert (new_cell.alpha_eff, new_cell.beta_eff) == (1.0, 1.0) + + @pytest.mark.asyncio async def test_load_state_from_db_handles_unknown_request_type(): r = _make_router() diff --git a/tests/test_litellm/router_strategy/adaptive_router/test_bandit.py b/tests/test_litellm/router_strategy/adaptive_router/test_bandit.py index ab322f0fb37..d773e56004f 100644 --- a/tests/test_litellm/router_strategy/adaptive_router/test_bandit.py +++ b/tests/test_litellm/router_strategy/adaptive_router/test_bandit.py @@ -5,6 +5,7 @@ import pytest from litellm.router_strategy.adaptive_router.bandit import ( BanditCell, apply_delta, + apply_efficiency_delta, initial_cell, normalized_cost, pick_best, @@ -67,6 +68,27 @@ def test_apply_delta_respects_sample_cap(): assert same_cell.beta == cell.beta +def test_apply_efficiency_delta_increments_alpha_eff_and_beta_eff(): + cell = BanditCell(alpha=5.0, beta=5.0, alpha_eff=1.0, beta_eff=1.0) + new_cell = apply_efficiency_delta(cell, 0.8, 0.2) + assert new_cell.alpha_eff == 1.8 + assert new_cell.beta_eff == 1.2 + assert new_cell.alpha == 5.0 + assert new_cell.beta == 5.0 + + +def test_apply_efficiency_delta_respects_sample_cap(): + cell = BanditCell(alpha=5.0, beta=5.0, alpha_eff=SAMPLE_CAP - 1.0, beta_eff=1.0) + same_cell = apply_efficiency_delta(cell, 5.0, 5.0) + assert same_cell.alpha_eff == cell.alpha_eff + assert same_cell.beta_eff == cell.beta_eff + + +def test_efficiency_mean(): + cell = BanditCell(alpha=5.0, beta=5.0, alpha_eff=6.0, beta_eff=4.0) + assert abs(cell.efficiency_mean - 0.6) < 0.001 + + def test_thompson_sample_in_range(): cell = BanditCell(alpha=10.0, beta=5.0) rng = random.Random(42) diff --git a/tests/test_litellm/router_strategy/adaptive_router/test_state_endpoint.py b/tests/test_litellm/router_strategy/adaptive_router/test_state_endpoint.py index 5662870a5cb..d6de9ed4d6b 100644 --- a/tests/test_litellm/router_strategy/adaptive_router/test_state_endpoint.py +++ b/tests/test_litellm/router_strategy/adaptive_router/test_state_endpoint.py @@ -49,7 +49,7 @@ async def test_get_state_snapshot_returns_cell_per_request_type_per_model(): # Top-level shape assert snap["router_name"] == "r1" assert snap["available_models"] == ["fast", "smart"] - assert snap["weights"] == {"quality": 0.7, "cost": 0.3} + assert snap["weights"] == {"quality": 0.7, "cost": 0.3, "efficiency": 0.0} assert snap["model_costs"] == {"fast": 0.0001, "smart": 0.001} assert set(snap["queue"].keys()) == { "state_pending", @@ -66,8 +66,11 @@ async def test_get_state_snapshot_returns_cell_per_request_type_per_model(): "model", "alpha", "beta", + "alpha_eff", + "beta_eff", "samples", "quality_mean", + "efficiency_mean", } assert cell["model"] in {"fast", "smart"} assert cell["request_type"] in {rt.value for rt in RequestType} diff --git a/tests/test_litellm/router_strategy/adaptive_router/test_update_queue.py b/tests/test_litellm/router_strategy/adaptive_router/test_update_queue.py index 9baa69a19e0..77a608c784c 100644 --- a/tests/test_litellm/router_strategy/adaptive_router/test_update_queue.py +++ b/tests/test_litellm/router_strategy/adaptive_router/test_update_queue.py @@ -80,6 +80,31 @@ async def test_flush_state_sums_correctly(queue, mock_prisma): assert call.kwargs["data"]["create"]["total_samples"] == 2 +@pytest.mark.asyncio +async def test_add_state_delta_with_efficiency_deltas(queue): + """Efficiency deltas should aggregate separately from quality deltas.""" + await queue.add_state_delta( + "r1", "general", "gpt-4", delta_alpha=1.0, delta_beta=0.0, delta_alpha_eff=0.8, delta_beta_eff=0.2 + ) + sizes = await queue.queue_size() + assert sizes["state_pending"] == 1 + + +@pytest.mark.asyncio +async def test_flush_state_writes_efficiency_deltas(queue, mock_prisma): + """First-write creates include the Beta(1,1) efficiency prior plus the delta, so a + freshly-created row's efficiency posterior matches what's already in memory instead + of losing its prior on restart.""" + await queue.add_state_delta("r1", "general", "gpt-4", 1.0, 0.0, 0.8, 0.2) + await queue.flush_state_to_db(mock_prisma) + + call = mock_prisma.db.litellm_adaptiverouterstate.upsert.call_args + assert call.kwargs["data"]["create"]["alpha_eff"] == 1.8 + assert call.kwargs["data"]["create"]["beta_eff"] == 1.2 + assert "alpha_eff" in call.kwargs["data"]["update"] + assert "beta_eff" in call.kwargs["data"]["update"] + + @pytest.mark.asyncio async def test_flush_session_drains_aggregator(queue, mock_prisma): await queue.add_session_state("s1", "r1", "gpt-4", {"classified_type": "general"})