From 04f9b771bc8d131b86ea39672fb27a437efdbe58 Mon Sep 17 00:00:00 2001 From: Tin Date: Tue, 1 Sep 2026 20:27:52 -0700 Subject: [PATCH] feat(router): seed adaptive routing from evaluations --- .../router_strategy/adaptive_router/README.md | 37 +++++++-- .../adaptive_router/adaptive_router.py | 36 ++++++++- .../router_strategy/adaptive_router/bandit.py | 24 +++++- .../adaptive_router/training.py | 75 +++++++++++++++++++ litellm/types/router.py | 26 +++++++ .../adaptive_router/test_adaptive_router.py | 32 +++++++- .../adaptive_router/test_config.py | 43 +++++++++++ .../adaptive_router/test_training.py | 71 ++++++++++++++++++ 8 files changed, 330 insertions(+), 14 deletions(-) create mode 100644 litellm/router_strategy/adaptive_router/training.py create mode 100644 tests/test_litellm/router_strategy/adaptive_router/test_training.py diff --git a/litellm/router_strategy/adaptive_router/README.md b/litellm/router_strategy/adaptive_router/README.md index 09420a8dd9d..6d21339a89e 100644 --- a/litellm/router_strategy/adaptive_router/README.md +++ b/litellm/router_strategy/adaptive_router/README.md @@ -39,6 +39,7 @@ model_list: adaptive_router_default_model: gpt-4o-mini adaptive_router_config: available_models: ["gpt-4o", "gpt-4o-mini"] + exploration_rate: 0.05 weights: quality: 0.7 cost: 0.3 @@ -47,15 +48,41 @@ model_list: Callers may pass header `x-litellm-min-quality-tier: 3` (or metadata key `min_quality_tier: 3`) to force selection from tier-3-or-higher models only. +## Seed quality from offline evaluations + +The router can start from measured quality instead of relying only on declared +tiers and strengths. Store one JSON object per evaluated model response: + +```json +{"request_type":"code_generation","model":"gpt-4o-mini","quality":1.0} +{"request_type":"code_generation","model":"gpt-4o","quality":0.8} +``` + +`quality` accepts partial-credit values from 0 through 1. Generate aggregated +Beta priors with: + +```shell +python -m litellm.router_strategy.adaptive_router.training evaluations.jsonl +``` + +Add the resulting `evaluation_priors` to `adaptive_router_config`. The router +combines each prior with its declared quality tier and strength. Large datasets +are compressed to `evaluation_prior_max_mass`, which defaults to 50, so online +feedback can still change the posterior after deployment. The generated config +uses posterior means for 95% of requests and Thompson sampling for 5% so the +router continues to explore without paying the cost on every request. + ## Behavior summary - **Cold start.** Each `(request_type, model)` cell starts with a Beta prior whose mean = `BASE_TIER_WEIGHT[tier] (+ STRENGTH_BONUS if declared)` - 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. + and total mass = `COLD_START_MASS` (10). Configured `evaluation_priors` update + and optionally strengthen that prior before online observations arrive. +- **Per-request decision.** Use each model's posterior mean for exploitation + requests and a Thompson sample for exploration requests. Score with + `quality_weight·estimate + cost_weight·normalized_cost`, then pick the argmax. + `exploration_rate` defaults to 1 for backward compatibility. Offline training + emits the recommended value of 0.05. - **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 diff --git a/litellm/router_strategy/adaptive_router/adaptive_router.py b/litellm/router_strategy/adaptive_router/adaptive_router.py index 1a33ea23bd4..cd9b2c071a5 100644 --- a/litellm/router_strategy/adaptive_router/adaptive_router.py +++ b/litellm/router_strategy/adaptive_router/adaptive_router.py @@ -16,7 +16,9 @@ from __future__ import annotations import asyncio import time from collections import OrderedDict +from collections.abc import Mapping from dataclasses import asdict, dataclass +from types import MappingProxyType from typing import Any, Final, cast from litellm._logging import verbose_router_logger @@ -26,6 +28,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( from litellm.router_strategy.adaptive_router.bandit import ( BanditCell, apply_delta, + apply_evaluation_prior, initial_cell, pick_best, ) @@ -60,6 +63,7 @@ from litellm.repositories.table_repositories import AdaptiveRouterStateRepositor from litellm.types.llms.openai import AllMessageValues from litellm.types.router import ( AdaptiveRouterConfig, + AdaptiveRouterEvaluationPrior, AdaptiveRouterPreferences, PreRoutingHookResponse, RequestType, @@ -117,10 +121,32 @@ class AdaptiveRouter: def _init_cold_start_cells(self) -> None: """Populate _cells with cold-start priors for every (rt, model) combination.""" - for rt in RequestType: - for model in self.config.available_models: - prefs = self.model_to_prefs.get(model) or _default_prefs() - self._cells[(rt, model)] = initial_cell(prefs, rt) + configured_priors: Final = MappingProxyType( + {(prior.request_type, prior.model): prior for prior in self.config.evaluation_priors} + ) + self._cells = { # mutable-ok: online feedback updates bandit cells in place + (request_type, model): self._initial_cell(request_type, model, configured_priors) + for request_type in RequestType + for model in self.config.available_models + } + + def _initial_cell( + self, + request_type: RequestType, + model: str, + configured_priors: Mapping[tuple[RequestType, str], AdaptiveRouterEvaluationPrior], + ) -> BanditCell: + prefs: Final = self.model_to_prefs.get(model) or _default_prefs() + base_cell: Final = initial_cell(prefs, request_type) + prior: Final = configured_priors.get((request_type, model)) + if prior is None: + return base_cell + return apply_evaluation_prior( + cell=base_cell, + successes=prior.successes, + failures=prior.failures, + max_mass=self.config.evaluation_prior_max_mass, + ) async def load_state_from_db(self, prisma_client: Any) -> None: """Override cold-start cells with persisted state. Called once at startup.""" @@ -225,6 +251,7 @@ class AdaptiveRouter: costs, quality_weight=self.config.weights.quality, cost_weight=self.config.weights.cost, + exploration_rate=self.config.exploration_rate, ) async def get_state_snapshot(self) -> dict[str, Any]: @@ -256,6 +283,7 @@ class AdaptiveRouter: "quality": self.config.weights.quality, "cost": self.config.weights.cost, }, + "exploration_rate": self.config.exploration_rate, "model_costs": dict(self.model_to_cost), "cells": cells, "feedback_contexts_live": feedback_contexts_live, diff --git a/litellm/router_strategy/adaptive_router/bandit.py b/litellm/router_strategy/adaptive_router/bandit.py index c2eb232bd99..07b904abfba 100644 --- a/litellm/router_strategy/adaptive_router/bandit.py +++ b/litellm/router_strategy/adaptive_router/bandit.py @@ -7,7 +7,7 @@ Each (router, request_type, model) cell is a Beta(alpha, beta) posterior. - mean = alpha / (alpha + beta) - total samples = alpha + beta - COLD_START_MASS (informative prior, not data) -Hot path: thompson_sample() — pure function, no I/O. +Hot path: posterior mean or thompson_sample() — pure functions, no I/O. """ import random @@ -75,6 +75,21 @@ def apply_delta(cell: BanditCell, delta_alpha: float, delta_beta: float) -> Band return BanditCell(alpha=new_alpha, beta=new_beta) +def apply_evaluation_prior( + cell: BanditCell, + successes: float, + failures: float, + max_mass: float, +) -> BanditCell: + alpha: Final = cell.alpha + successes + beta: Final = cell.beta + failures + total: Final = alpha + beta + if total <= max_mass: + return BanditCell(alpha=alpha, beta=beta) + mean: Final = alpha / total + return BanditCell(alpha=mean * max_mass, beta=(1.0 - mean) * max_mass) + + 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].""" r: Final = rng if rng is not None else random @@ -114,10 +129,11 @@ def pick_best( model_costs: dict[str, float], quality_weight: float = DEFAULT_QUALITY_WEIGHT, cost_weight: float = DEFAULT_COST_WEIGHT, + exploration_rate: float = 1.0, rng: random.Random | None = None, ) -> str: """ - Sample once per model, score each, return the model with highest score. + Use posterior means for exploitation or sample once per model for exploration. cells: {model_name: BanditCell} model_costs: {model_name: $/1k tokens} @@ -125,10 +141,12 @@ def pick_best( if not cells: raise ValueError("pick_best called with no models") all_costs: Final = list(model_costs.values()) + random_source: Final = rng if rng is not None else random + explore: Final = random_source.random() < exploration_rate best_model: str | None = None best_score = float("-inf") for model, cell in cells.items(): - q = thompson_sample(cell, rng=rng) + q = thompson_sample(cell, rng=rng) if explore else cell.mean s = score(q, model_costs[model], all_costs, quality_weight, cost_weight) if s > best_score: best_score = s diff --git a/litellm/router_strategy/adaptive_router/training.py b/litellm/router_strategy/adaptive_router/training.py new file mode 100644 index 00000000000..008d4d840d4 --- /dev/null +++ b/litellm/router_strategy/adaptive_router/training.py @@ -0,0 +1,75 @@ +from __future__ import annotations + +import sys +from collections.abc import Iterable, Iterator +from itertools import groupby +from pathlib import Path +from typing import Final, TypeAlias + +from pydantic import BaseModel, ConfigDict, Field + +from litellm.types.router import ( + AdaptiveRouterEvaluationPrior, + RequestType, +) + +EvaluationKey: TypeAlias = tuple[RequestType, str] + + +class AdaptiveRouterEvaluationRecord(BaseModel): + model_config = ConfigDict(frozen=True) + + request_type: RequestType + model: str + quality: float = Field(ge=0.0, le=1.0) + + +class AdaptiveRouterTrainingResult(BaseModel): + model_config = ConfigDict(frozen=True) + + evaluation_priors: tuple[AdaptiveRouterEvaluationPrior, ...] + exploration_rate: float = 0.05 + + +def _aggregate_group( + key: EvaluationKey, + records: Iterator[AdaptiveRouterEvaluationRecord], +) -> AdaptiveRouterEvaluationPrior: + request_type, model = key + grouped_records: Final = tuple(records) + return AdaptiveRouterEvaluationPrior( + request_type=request_type, + model=model, + successes=sum(record.quality for record in grouped_records), + failures=sum(1.0 - record.quality for record in grouped_records), + ) + + +def aggregate_evaluation_records( + records: Iterable[AdaptiveRouterEvaluationRecord], +) -> tuple[AdaptiveRouterEvaluationPrior, ...]: + ordered: Final = tuple(sorted(records, key=lambda record: (record.request_type.value, record.model))) + grouped: Final = groupby(ordered, key=lambda record: (record.request_type, record.model)) + return tuple(_aggregate_group(key, records_for_key) for key, records_for_key in grouped) + + +def load_evaluation_records(path: Path) -> tuple[AdaptiveRouterEvaluationRecord, ...]: + with path.open(encoding="utf-8") as input_file: + return tuple(AdaptiveRouterEvaluationRecord.model_validate_json(line) for line in input_file if line.strip()) + + +def training_config_fragment(path: Path) -> AdaptiveRouterTrainingResult: + priors: Final = aggregate_evaluation_records(load_evaluation_records(path)) + return AdaptiveRouterTrainingResult(evaluation_priors=priors) + + +def main() -> None: + if len(sys.argv) != 2: + raise SystemExit("usage: python -m litellm.router_strategy.adaptive_router.training EVALUATIONS.jsonl") + evaluation_records: Final = Path(sys.argv[1]) + result: Final = training_config_fragment(evaluation_records) + sys.stdout.write(result.model_dump_json(indent=2) + "\n") + + +if __name__ == "__main__": + main() diff --git a/litellm/types/router.py b/litellm/types/router.py index e0957383aac..c99df275821 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -1036,9 +1036,35 @@ class AdaptiveRouterWeights(BaseModel): return v +class AdaptiveRouterEvaluationPrior(BaseModel): + request_type: RequestType + model: str + successes: float = Field(ge=0.0) + failures: float = Field(ge=0.0) + + @model_validator(mode="after") + def _has_observations(self) -> "AdaptiveRouterEvaluationPrior": + if self.successes + self.failures <= 0: + raise ValueError("evaluation prior must contain at least one observation") + return self + + class AdaptiveRouterConfig(BaseModel): available_models: list[str] weights: AdaptiveRouterWeights = Field(default_factory=AdaptiveRouterWeights) + evaluation_priors: tuple[AdaptiveRouterEvaluationPrior, ...] = () + evaluation_prior_max_mass: float = Field(default=50.0, ge=10.0, le=100.0) + exploration_rate: float = Field(default=1.0, ge=0.0, le=1.0) + + @model_validator(mode="after") + def _valid_evaluation_priors(self) -> "AdaptiveRouterConfig": + keys: Final = tuple((prior.request_type, prior.model) for prior in self.evaluation_priors) + unknown_models: Final = sorted({model for _, model in keys if model not in self.available_models}) + if unknown_models: + raise ValueError(f"evaluation priors reference unavailable models: {unknown_models}") + if len(keys) != len(set(keys)): + raise ValueError("evaluation priors must contain unique request_type and model pairs") + return self class AdaptiveRouterPreferences(BaseModel): 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..2de3bc6b04f 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 @@ -1,15 +1,16 @@ """Unit tests for the AdaptiveRouter strategy class.""" +from typing import Final from unittest.mock import AsyncMock, MagicMock -from litellm.router_strategy.adaptive_router import adaptive_router as ar_module - import pytest +from litellm.router_strategy.adaptive_router import adaptive_router as ar_module from litellm.router_strategy.adaptive_router.adaptive_router import AdaptiveRouter from litellm.router_strategy.adaptive_router.signals import Turn from litellm.types.router import ( AdaptiveRouterConfig, + AdaptiveRouterEvaluationPrior, AdaptiveRouterPreferences, RequestType, ) @@ -37,6 +38,33 @@ async def test_pick_model_returns_model_from_available_list(): assert chosen in {"fast", "smart"} +@pytest.mark.asyncio +async def test_evaluation_priors_seed_matching_request_type_and_model() -> None: + cfg: Final = AdaptiveRouterConfig( + available_models=["fast", "smart"], + exploration_rate=0, + evaluation_priors=( + AdaptiveRouterEvaluationPrior( + request_type=RequestType.CODE_GENERATION, + model="fast", + successes=18, + failures=2, + ), + ), + ) + router: Final = AdaptiveRouter( + router_name="seeded", + config=cfg, + model_to_prefs={ + "fast": AdaptiveRouterPreferences(quality_tier=2), + "smart": AdaptiveRouterPreferences(quality_tier=2), + }, + model_to_cost={"fast": 0.001, "smart": 0.001}, + ) + + assert await router.pick_model(RequestType.CODE_GENERATION) == "fast" + + @pytest.mark.asyncio async def test_pick_model_min_quality_tier_filter(): r = _make_router() diff --git a/tests/test_litellm/router_strategy/adaptive_router/test_config.py b/tests/test_litellm/router_strategy/adaptive_router/test_config.py index fd14556a0bc..8f4dae00644 100644 --- a/tests/test_litellm/router_strategy/adaptive_router/test_config.py +++ b/tests/test_litellm/router_strategy/adaptive_router/test_config.py @@ -3,6 +3,7 @@ from pydantic import ValidationError from litellm.types.router import ( AdaptiveRouterConfig, + AdaptiveRouterEvaluationPrior, AdaptiveRouterPreferences, AdaptiveRouterWeights, # noqa: F401 # imported per spec, exercised transitively RequestType, @@ -53,3 +54,45 @@ def test_config_accepts_all_six_request_types_in_strengths(): ], ) assert len(prefs.strengths) == 6 + + +def test_config_accepts_evaluation_priors(): + cfg = AdaptiveRouterConfig( + available_models=["fast", "smart"], + evaluation_priors=( + AdaptiveRouterEvaluationPrior( + request_type=RequestType.CODE_GENERATION, + model="fast", + successes=80, + failures=20, + ), + ), + ) + + assert cfg.evaluation_priors[0].successes == 80 + + +def test_config_accepts_bounded_exploration_rate(): + cfg = AdaptiveRouterConfig(available_models=["fast", "smart"], exploration_rate=0.05) + + assert cfg.exploration_rate == 0.05 + + with pytest.raises(ValidationError): + AdaptiveRouterConfig(available_models=["fast", "smart"], exploration_rate=1.1) + + +def test_config_rejects_evaluation_prior_for_unavailable_model(): + with pytest.raises(ValidationError, match="unavailable models"): + AdaptiveRouterConfig.model_validate( + { + "available_models": ["fast"], + "evaluation_priors": [ + { + "request_type": "code_generation", + "model": "smart", + "successes": 1, + "failures": 1, + } + ], + } + ) diff --git a/tests/test_litellm/router_strategy/adaptive_router/test_training.py b/tests/test_litellm/router_strategy/adaptive_router/test_training.py new file mode 100644 index 00000000000..974495b20eb --- /dev/null +++ b/tests/test_litellm/router_strategy/adaptive_router/test_training.py @@ -0,0 +1,71 @@ +import random +from pathlib import Path +from typing import Final + +from litellm.router_strategy.adaptive_router.bandit import ( + BanditCell, + apply_evaluation_prior, + pick_best, +) +from litellm.router_strategy.adaptive_router.training import training_config_fragment + + +def test_apply_evaluation_prior_preserves_quality_and_caps_mass() -> None: + cell: Final = apply_evaluation_prior( + BanditCell(alpha=5.0, beta=5.0), + successes=90.0, + failures=10.0, + max_mass=50.0, + ) + + assert abs(cell.alpha + cell.beta - 50.0) < 0.001 + assert abs(cell.mean - (95.0 / 110.0)) < 0.001 + + +def test_pick_best_uses_posterior_mean_without_exploration() -> None: + chosen: Final = pick_best( + cells={ + "reliable": BanditCell(alpha=90.0, beta=10.0), + "unreliable": BanditCell(alpha=10.0, beta=90.0), + }, + model_costs={"reliable": 1.0, "unreliable": 1.0}, + exploration_rate=0.0, + rng=random.Random(42), + ) + + assert chosen == "reliable" + + +def test_training_config_fragment_aggregates_quality(tmp_path: Path) -> None: + evaluation_records: Final = tmp_path / "evaluation.jsonl" + evaluation_records.write_text( + "\n".join( + ( + '{"request_type":"code_generation","model":"fast","quality":1}', + '{"request_type":"code_generation","model":"fast","quality":0.5}', + '{"request_type":"code_generation","model":"fast","quality":0}', + '{"request_type":"writing","model":"smart","quality":1}', + ) + ), + encoding="utf-8", + ) + + fragment: Final = training_config_fragment(evaluation_records).model_dump(mode="json") + + assert fragment == { + "evaluation_priors": [ + { + "request_type": "code_generation", + "model": "fast", + "successes": 1.5, + "failures": 1.5, + }, + { + "request_type": "writing", + "model": "smart", + "successes": 1.0, + "failures": 0.0, + }, + ], + "exploration_rate": 0.05, + }