From 9fd60e4f95228718ed74f75ace470016ee4f42e4 Mon Sep 17 00:00:00 2001 From: tin-berri Date: Sat, 5 Sep 2026 19:24:00 -0700 Subject: [PATCH] feat(router): gate heuristic v1 tuning (#39952) --- .../model_management_endpoints.py | 95 ++++++--- litellm/proxy/proxy_server.py | 90 ++++++++- litellm/repositories/model_repository.py | 10 +- .../auto_router_tuning_baseline.py | 157 +++++++++++++++ .../test_model_management_endpoints.py | 186 +++++++++++++++++- .../proxy/proxy_server/test_lifecycle.py | 18 ++ .../test_auto_router_tuning_baseline.py | 185 +++++++++++++++++ 7 files changed, 710 insertions(+), 31 deletions(-) create mode 100644 litellm/router_utils/auto_router_tuning_baseline.py create mode 100644 tests/test_litellm/router_utils/test_auto_router_tuning_baseline.py diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index b77108911aa..0f19e9ce149 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -61,6 +61,7 @@ from litellm.proxy.common_utils.encrypt_decrypt_utils import ( encrypt_value_helper, ) from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.db.routing_prisma_wrapper import WriterPinnedClient from litellm.proxy.management_endpoints.common_utils import _is_user_team_admin from litellm.proxy.management_endpoints.team_endpoints import ( _refresh_cached_team, @@ -109,6 +110,7 @@ from litellm.router_utils.auto_router_model_naming import ( validate_complexity_router_config_write, validate_strategy_router_model_write, ) +from litellm.router_utils.auto_router_tuning_baseline import is_mutable_tuned_candidate, tuning_quota_violation from litellm.types.proxy.management_endpoints.model_management_endpoints import ( AutoRouterClassifierDefaultPromptResponse, UpdateUsefulLinksRequest, @@ -349,6 +351,36 @@ def _effective_complexity_router_params( ) +def _decrypted_model(stored_model: object) -> str | None: + if not isinstance(stored_model, str): + return None + decrypted: Final = decrypt_value_helper( + value=stored_model, key="model", exception_type="debug", return_original_value=True + ) + return decrypted if isinstance(decrypted, str) else None + + +def _tuning_candidate(effective_params: Mapping[str, object], model_id: str | None) -> Mapping[str, object]: + return MappingProxyType( + { + "litellm_params": effective_params, + "model_info": MappingProxyType({"id": model_id, "db_model": True}), + } + ) + + +def _raise_on_tuning_quota_violation( + *, + candidate: Mapping[str, object], + others: Sequence[Mapping[str, object]], + baselines: Mapping[str, str], + limit: int | None, +) -> None: + violation: Final = tuning_quota_violation(candidate=candidate, others=others, baselines=baselines, limit=limit) + if violation is not None: + raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail=f"{violation} {AUTO_ROUTER_LICENSE_REMEDY}") + + @asynccontextmanager async def _auto_router_capability_slot( prisma_client: PrismaClient, *, effective_params: Mapping[str, object], model_id: str | None @@ -366,40 +398,55 @@ async def _auto_router_capability_slot( must wait until the transaction has committed and the lock is released. The transaction writes bypass the repository's publish-on-write, so the config change is published once after commit, the way delete_team_models does. + + A heuristic-v1 router whose tuning has moved off its recorded baseline is judged the same + way under the same lock, against the DB rows plus this proxy's config.yaml routers. """ - from litellm.proxy.proxy_server import _license_check, llm_router + from litellm.proxy.proxy_server import ( + _license_check, # pyright: ignore[reportPrivateUsage] # existing capability slot reads the proxy license singleton + heuristic_v1_tuning_baselines, + llm_router, + ) limit: Final = _license_check.auto_router_capability_limit() capability: Final = gated_capability_of(effective_params) - if limit is None or capability is None: + baselines: Final = heuristic_v1_tuning_baselines + tuning_candidate: Final = _tuning_candidate(effective_params, model_id=model_id) + judges_tuning: Final = baselines is not None and is_mutable_tuned_candidate(tuning_candidate, baselines) + if limit is None or (capability is None and not judges_tuning): yield _proxy_model_table(prisma_client) return async with prisma_client.db.tx() as tx_ctx: tables: Final[_TxModelTables] = tx_ctx await tx_ctx.query_raw(_CAPABILITY_LOCK_SQL, AUTO_ROUTER_CAPABILITY_SLOT_LOCK_KEY) - rows: Sequence[Mapping[str, object]] = await tx_ctx.query_raw( - _CAPABILITY_DB_ROWS_SQL[capability.key], model_id or "" - ) - db_held: Final = sum( - 1 - for row in rows - for stored_model in (row.get("model"),) - if isinstance(stored_model, str) - and is_complexity_router_model( - decrypt_value_helper( - value=stored_model, - key="model", - exception_type="debug", - return_original_value=True, - ) - ) - ) config_rows: Final = () if llm_router is None else tuple(llm_router.config_deployments()) - held: Final = db_held + count_capability_routers(config_rows, capability=capability) - violation: Final = capability_limit_violation(capability=capability, held=held + 1, limit=limit) - if violation is not None: - raise HTTPException( - status_code=status.HTTP_403_FORBIDDEN, detail=f"{violation} {AUTO_ROUTER_LICENSE_REMEDY}" + if capability is not None: + rows: Sequence[Mapping[str, object]] = await tx_ctx.query_raw( + _CAPABILITY_DB_ROWS_SQL[capability.key], model_id or "" + ) + db_held: Final = sum(1 for row in rows if is_complexity_router_model(_decrypted_model(row.get("model")))) + held: Final = db_held + count_capability_routers(config_rows, capability=capability) + violation: Final = capability_limit_violation(capability=capability, held=held + 1, limit=limit) + if violation is not None: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, detail=f"{violation} {AUTO_ROUTER_LICENSE_REMEDY}" + ) + if judges_tuning and baselines is not None: + model_rows: Final = await ModelRepository(WriterPinnedClient(tx_ctx)).find_all_except(model_id or "") + _raise_on_tuning_quota_violation( + candidate=tuning_candidate, + others=tuple( + MappingProxyType( + { + "litellm_params": row.litellm_params, + "model_info": MappingProxyType({"id": row.model_id, "db_model": True}), + } + ) + for row in model_rows + ) + + config_rows, + baselines=baselines, + limit=limit, ) yield tables.litellm_proxymodeltable await publish_config_change(redis_cache=coordination_redis_cache(), object_type="litellm_proxymodeltable") diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 1cd08fe27c0..ba5714fe950 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -125,6 +125,12 @@ from litellm.router_utils.auto_router_model_naming import ( count_capability_routers, validate_complexity_router_config_placement, ) +from litellm.router_utils.auto_router_tuning_baseline import ( + TUNING_BASELINE_PARAM_NAME, + mutable_tuned_identities, + snapshot_tuning_baselines, + tuning_limit_violation, +) from litellm.types.utils import ( ModelResponse, ModelResponseStream, @@ -892,7 +898,8 @@ def cleanup_router_config_variables(): use_shared_health_check, \ health_check_interval, \ health_check_concurrency, \ - prisma_client + prisma_client, \ + heuristic_v1_tuning_baselines # Set all variables to None master_key = None @@ -911,6 +918,7 @@ def cleanup_router_config_variables(): health_check_interval = None health_check_concurrency = None prisma_client = None + heuristic_v1_tuning_baselines = None async def _flush_spend_logs_queue_on_shutdown() -> None: @@ -2267,6 +2275,7 @@ experimental = False #### GLOBAL VARIABLES #### llm_router: Router | None = None llm_model_list: list | None = None +heuristic_v1_tuning_baselines: Mapping[str, str] | None = None # Serializes every model reconcile (ProxyConfig.add_deployment and clear_cache) so the # read-modify-write of llm_router above is atomic. Without it, two concurrent model # writes each reconcile the router against their OWN db snapshot, and the one holding @@ -9324,6 +9333,77 @@ class ProxyStartupEvent: except Exception as e: verbose_proxy_logger.debug("UI settings sync on startup skipped or failed: %s", e) + @classmethod + async def _load_heuristic_v1_tuning_baselines( + cls, prisma_client: PrismaClient, deployments: Sequence[Mapping[str, object]] + ) -> Mapping[str, str] | None: + """Read the recorded tuning baselines, recording current routers on the first boot.""" + from prisma.errors import UniqueViolationError + + try: + config_table: Final = prisma_client.db.litellm_config + row: Final = await config_table.find_unique( + where={"param_name": TUNING_BASELINE_PARAM_NAME} # mutable-ok: Prisma rejects mappingproxy input + ) + if row is not None: + stored: Final = row.param_value + decoded: Final = json.loads(stored) if isinstance(stored, str) else stored + return MappingProxyType( + { + str(identity): str(fingerprint) + for identity, fingerprint in (decoded.items() if isinstance(decoded, Mapping) else ()) + } + ) # mutable-ok: MappingProxyType owns the completed immutable baseline + snapshot: Final = snapshot_tuning_baselines(deployments) + try: + await config_table.create( + data={ # mutable-ok: Prisma rejects mappingproxy input + "param_name": TUNING_BASELINE_PARAM_NAME, + "param_value": json.dumps(dict(snapshot)), # mutable-ok: json only serializes concrete mappings + } + ) + verbose_proxy_logger.info("Recorded heuristic-v1 tuning baseline for %s auto-router(s)", len(snapshot)) + return snapshot + except UniqueViolationError: + competing_row: Final = await config_table.find_unique( + where={"param_name": TUNING_BASELINE_PARAM_NAME} # mutable-ok: Prisma rejects mappingproxy input + ) + competing_value: Final = None if competing_row is None else competing_row.param_value + competing_decoded: Final = ( + json.loads(competing_value) if isinstance(competing_value, str) else competing_value + ) + return MappingProxyType( + { + str(identity): str(fingerprint) + for identity, fingerprint in ( + competing_decoded.items() if isinstance(competing_decoded, Mapping) else () + ) + } + ) # mutable-ok: MappingProxyType owns the completed immutable baseline + except Exception as e: # noqa: BLE001 # enforcement is skipped for this boot; refusing every tuned router on a DB blip is the one outcome the gate forbids + verbose_proxy_logger.warning("Heuristic-v1 tuning baseline unavailable, gate not enforced this boot: %s", e) + return None + + @classmethod + async def enforce_heuristic_v1_tuning_baseline( + cls, prisma_client: PrismaClient, llm_router: Router | None, limit: int | None + ) -> Mapping[str, str] | None: + """Load a complete baseline and reject a startup that exceeds the tuning quota.""" + db_models: Final = await proxy_config._get_models_from_db(prisma_client) + if db_models is None: + verbose_proxy_logger.warning("Heuristic-v1 tuning baseline unavailable, gate not enforced this boot") + return None + config_deployments: Final = () if llm_router is None else tuple(llm_router.config_deployments()) + deployments: Final = (*config_deployments, *proxy_config.decrypt_model_list_from_db(db_models)) + baselines: Final = await cls._load_heuristic_v1_tuning_baselines(prisma_client, deployments) + if baselines is None: + return None + mutable: Final = mutable_tuned_identities(deployments, baselines) + violation: Final = tuning_limit_violation(held=len(mutable), limit=limit) + if violation is not None: + raise ValueError(f"model_list: {violation} {AUTO_ROUTER_LICENSE_REMEDY}") + return baselines + @classmethod async def initialize_scheduled_background_jobs( cls, @@ -9335,7 +9415,7 @@ class ProxyStartupEvent: proxy_logging_obj: ProxyLogging, ) -> ProxyWorkerHeartbeat: """Initializes scheduled background jobs""" - global store_model_in_db, scheduler + global heuristic_v1_tuning_baselines, store_model_in_db, scheduler # rebind-ok: startup publishes the one read-only baseline snapshot # MEMORY LEAK FIX: Configure scheduler with optimized settings # Memray analysis showed APScheduler's normalize() and _apply_jitter() causing @@ -9573,6 +9653,12 @@ class ProxyStartupEvent: misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME, ) + heuristic_v1_tuning_baselines = await cls.enforce_heuristic_v1_tuning_baseline( + prisma_client=prisma_client, + llm_router=llm_router, + limit=_license_check.auto_router_capability_limit(), + ) + await cls._initialize_slack_alerting_jobs( scheduler=scheduler, general_settings=general_settings, diff --git a/litellm/repositories/model_repository.py b/litellm/repositories/model_repository.py index acc7c8dcda8..540f197044f 100644 --- a/litellm/repositories/model_repository.py +++ b/litellm/repositories/model_repository.py @@ -3,7 +3,8 @@ Model repository for database operations on LiteLLM_ProxyModelTable. """ import json -from collections.abc import Mapping +from collections.abc import Mapping, Sequence +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Protocol from litellm.models.model import LiteLLM_ProxyModelTable @@ -105,6 +106,13 @@ class ModelRepository(BaseRepository[LiteLLM_ProxyModelTable]): records: Final = await self.table.find_many(where={"blocked": False}) return self._to_model_list(records) + async def find_all_except(self, model_id: str) -> Sequence[LiteLLM_ProxyModelTable]: + """Find every model except the row currently being updated.""" + records: Final = await self.table.find_many( + where=MappingProxyType({"model_id": MappingProxyType({"not": model_id})}) + ) + return tuple(self._to_model_list(records)) + async def find_by_team_id(self, team_id: str) -> list[LiteLLM_ProxyModelTable]: """Find models associated with a specific team. diff --git a/litellm/router_utils/auto_router_tuning_baseline.py b/litellm/router_utils/auto_router_tuning_baseline.py new file mode 100644 index 00000000000..1e6412856a6 --- /dev/null +++ b/litellm/router_utils/auto_router_tuning_baseline.py @@ -0,0 +1,157 @@ +"""Baseline-relative license gate for heuristic-v1 complexity-router tuning.""" + +import hashlib +import json +from collections.abc import Iterable, Mapping +from types import MappingProxyType +from typing import Final + +from pydantic import ValidationError + +from litellm.router_strategy.complexity_router.config import ComplexityRouterConfig + +TUNING_BASELINE_PARAM_NAME: Final = "auto_router_tuning_baseline" + +HEURISTIC_V1_TUNING_FIELDS: Final = ( + "tiers", + "tier_model_configs", + "classifier_type", + "tier_boundaries", + "reasoning_override_min_score", + "token_thresholds", + "dimension_weights", + "code_keywords", + "reasoning_keywords", + "technical_keywords", + "custom_technical_keywords", + "simple_keywords", + "escalation_keywords", + "keyword_tier_rules", +) + +_V1_SCORING_CLASSIFIER_TYPES: Final = frozenset({"heuristic", "heuristic_first", "hybrid"}) +_AUTO_ROUTER_COMPLEXITY_PREFIX: Final = "auto_router/complexity_router" +_EMPTY: Final[Mapping[str, object]] = MappingProxyType({}) +_EMPTY_TAGS: Final[tuple[str, ...]] = () + + +def _mapping(value: object) -> Mapping[str, object]: + return value if isinstance(value, Mapping) else _EMPTY + + +def tuning_fingerprint(complexity_router_config: object) -> str | None: + """Digest of normalized heuristic-v1 tuning fields, or None when the config is invalid.""" + try: + validated: Final = ComplexityRouterConfig.model_validate(_mapping(complexity_router_config)) + except ValidationError: + return None + payload: Final = validated.model_dump(mode="json", include=frozenset(HEURISTIC_V1_TUNING_FIELDS)) + return hashlib.sha256(json.dumps(payload, sort_keys=True, separators=(",", ":")).encode()).hexdigest() + + +DEFAULT_TUNING_FINGERPRINT: Final = tuning_fingerprint(_EMPTY) + + +def uses_heuristic_v1(complexity_router_config: object) -> bool: + """Whether a config's primary classifier path is the heuristic-v1 scorer.""" + return _mapping(complexity_router_config).get("classifier_type", "heuristic") in _V1_SCORING_CLASSIFIER_TYPES + + +def router_identity(deployment: Mapping[str, object]) -> str | None: + """Stable identity for a complexity-router deployment, across tuning edits.""" + model_info: Final = _mapping(deployment.get("model_info")) + model_id: Final = model_info.get("id") + if model_info.get("db_model") is True and isinstance(model_id, str) and model_id: + return f"db:{model_id}" + model_name: Final = deployment.get("model_name") + if not isinstance(model_name, str) or not model_name: + return None + litellm_params: Final = _mapping(deployment.get("litellm_params")) + tags: Final = litellm_params.get("tags") + normalized_tags: Final = ( + tuple(sorted(str(tag) for tag in tags)) + if isinstance(tags, Iterable) and not isinstance(tags, str) + else _EMPTY_TAGS + ) + return f"yaml:{json.dumps((model_name, normalized_tags), separators=(',', ':'))}" + + +def heuristic_v1_router_fingerprint(deployment: Mapping[str, object]) -> tuple[str, str] | None: + """The identity/fingerprint pair for a heuristic-v1 complexity router, else None.""" + litellm_params: Final = _mapping(deployment.get("litellm_params")) + model: Final = litellm_params.get("model") + config: Final = litellm_params.get("complexity_router_config") + if ( + not isinstance(model, str) + or not model.startswith(_AUTO_ROUTER_COMPLEXITY_PREFIX) + or not uses_heuristic_v1(config) + ): + return None + identity: Final = router_identity(deployment) + fingerprint: Final = tuning_fingerprint(config) + if identity is None or fingerprint is None: + return None + return identity, fingerprint + + +def snapshot_tuning_baselines(deployments: Iterable[Mapping[str, object]]) -> Mapping[str, str]: + """One immutable first-observation baseline for every heuristic-v1 complexity router.""" + return MappingProxyType( + { + identity: fingerprint + for deployment in deployments + if (pair := heuristic_v1_router_fingerprint(deployment)) is not None + for identity, fingerprint in (pair,) + } + ) # mutable-ok: MappingProxyType owns the completed immutable snapshot + + +def is_mutable_tuned_candidate(candidate: Mapping[str, object], baselines: Mapping[str, str]) -> bool: + pair: Final = heuristic_v1_router_fingerprint(candidate) + if pair is None: + return False + identity, fingerprint = pair + return fingerprint != baselines.get(identity, DEFAULT_TUNING_FINGERPRINT) + + +def mutable_tuned_identities( + deployments: Iterable[Mapping[str, object]], baselines: Mapping[str, str] +) -> frozenset[str]: + """Heuristic-v1 routers whose current tuning differs from their baseline or shipped default.""" + return frozenset( + identity + for deployment in deployments + if (pair := heuristic_v1_router_fingerprint(deployment)) is not None + for identity, _ in (pair,) + if is_mutable_tuned_candidate(deployment, baselines) + ) + + +def tuning_limit_violation(*, held: int, limit: int | None) -> str | None: + if limit is None or held <= limit: + return None + return ( + f"At most {limit} auto-router(s) with changed heuristic scorer settings or tier models can be modified " + "without an auto-router license. Keep this router on its recorded settings, or revert the other changed " + "router to its baseline, or remove one of them." + ) + + +def tuning_quota_violation( + *, + candidate: Mapping[str, object], + others: Iterable[Mapping[str, object]], + baselines: Mapping[str, str], + limit: int | None, +) -> str | None: + """Why a change to candidate tuning exceeds the baseline-relative free quota.""" + if limit is None: + return None + pair: Final = heuristic_v1_router_fingerprint(candidate) + if pair is None: + return None + identity, _ = pair + if not is_mutable_tuned_candidate(candidate, baselines): + return None + held: Final = mutable_tuned_identities(others, baselines) - frozenset((identity,)) + return tuning_limit_violation(held=len(held) + 1, limit=limit) diff --git a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py index 33de2a09626..63c08529b4d 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py @@ -4440,13 +4440,28 @@ class TestStrategyRouterWriteValidation: class _FakeTx: """Stands in for a prisma transaction: records raw statements and returns encrypted-model candidates.""" - def __init__(self, db_models: list[str]) -> None: + def __init__(self, db_models: list[str], tuning_rows: list[dict[str, object]] | None = None) -> None: self.db_models = db_models + self.tuning_rows = tuning_rows or [] self.raw_calls: list[tuple[str, tuple[object, ...]]] = [] - self.litellm_proxymodeltable = MagicMock(create=AsyncMock(), update=AsyncMock()) + self.litellm_proxymodeltable = MagicMock( + create=AsyncMock(), + update=AsyncMock(), + find_many=AsyncMock( + return_value=tuple( + LiteLLM_ProxyModelTable.model_validate(row) for row in self.tuning_rows + ) + ), + ) + + @property + def db(self) -> "TestStrategyRouterWriteValidation._FakeTx": + return self async def query_raw(self, sql: str, *args: object) -> list[dict[str, object]]: self.raw_calls.append((sql, args)) + if "AS litellm_params" in sql: + return [row for row in self.tuning_rows if row.get("model_id") != args[0]] return [{"model": model} for model in self.db_models] if "AS model" in sql else [] async def __aenter__(self) -> "TestStrategyRouterWriteValidation._FakeTx": @@ -4458,9 +4473,11 @@ class TestStrategyRouterWriteValidation: class _FakeDb: """Stands in for prisma_client: the plain client and the transaction it opens are told apart by identity.""" - def __init__(self, db_models: list[str], existing_row: object = None) -> None: + def __init__( + self, db_models: list[str], existing_row: object = None, tuning_rows: list[dict[str, object]] | None = None + ) -> None: self.db = self - self.tx_obj = TestStrategyRouterWriteValidation._FakeTx(db_models) + self.tx_obj = TestStrategyRouterWriteValidation._FakeTx(db_models, tuning_rows=tuning_rows) self.litellm_proxymodeltable = MagicMock( create=AsyncMock(), update=AsyncMock(), find_unique=AsyncMock(return_value=existing_row) ) @@ -4617,6 +4634,167 @@ class TestStrategyRouterWriteValidation: assert capability is not None assert capability.sql_config_predicate.split("{config}")[-1].strip() in count_sql + _TUNED_A = {"classifier_type": "heuristic", "tiers": {"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o"}} + _TUNED_A_EDITED = {**_TUNED_A, "dimension_weights": {"codePresence": 0.9}} + _TUNED_B = {"classifier_type": "heuristic", "tiers": {"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4.1"}} + _TUNED_B_EDITED = {**_TUNED_B, "tiers": {"SIMPLE": "gpt-4o", "MEDIUM": "gpt-4.1"}} + + @staticmethod + def _db_router_row(model_id: str, config: Mapping[str, object]) -> dict[str, object]: + return { + "model_name": f"router-{model_id}", + "litellm_params": {"model": "auto_router/complexity_router", "complexity_router_config": dict(config)}, + "model_info": {"id": model_id, "db_model": True}, + } + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "limit,baseline_rows,live_rows,candidate_id,candidate_config,expected", + [ + (1, ["a", "b"], {"a": "_TUNED_A", "b": "_TUNED_B"}, "a", "_TUNED_A", "allowed"), + (1, ["a", "b"], {"a": "_TUNED_A", "b": "_TUNED_B"}, "a", "_TUNED_A_EDITED", "allowed"), + (1, ["a", "b"], {"a": "_TUNED_A_EDITED", "b": "_TUNED_B"}, "a", "_TUNED_A_EDITED", "allowed"), + (1, ["a", "b"], {"a": "_TUNED_A_EDITED", "b": "_TUNED_B"}, "b", "_TUNED_B_EDITED", "refused"), + (1, ["a", "b"], {"a": "_TUNED_A_EDITED", "b": "_TUNED_B"}, "c", "_TUNED_B", "refused"), + (1, ["a", "b"], {"a": "_TUNED_A_EDITED", "b": "_TUNED_B"}, "b", "_TUNED_B", "allowed"), + (None, ["a", "b"], {"a": "_TUNED_A_EDITED", "b": "_TUNED_B"}, "b", "_TUNED_B_EDITED", "allowed"), + (1, [], {}, "c", "_TUNED_B", "allowed"), + ], + ) + async def test_slot_enforces_baseline_relative_tuning_quota( + self, + limit: int | None, + baseline_rows: list[str], + live_rows: Mapping[str, str], + candidate_id: str, + candidate_config: str, + expected: str, + ) -> None: + """Without the license, one router may move off its recorded tuning baseline and keep being edited; + a change to a second router, or a second new tuned router, is refused. Unchanged baselines and + reverts to baseline are never counted, and a license lifts every check.""" + from fastapi import HTTPException + + from litellm.proxy.management_endpoints.model_management_endpoints import _auto_router_capability_slot + from litellm.router_utils.auto_router_tuning_baseline import snapshot_tuning_baselines + + configs = { + "_TUNED_A": self._TUNED_A, + "_TUNED_A_EDITED": self._TUNED_A_EDITED, + "_TUNED_B": self._TUNED_B, + "_TUNED_B_EDITED": self._TUNED_B_EDITED, + } + baselines = snapshot_tuning_baselines( + [self._db_router_row(row_id, configs["_TUNED_A" if row_id == "a" else "_TUNED_B"]) for row_id in baseline_rows] + ) + effective_params = { + "model": "auto_router/complexity_router", + "complexity_router_config": configs[candidate_config], + } + # The other routers live in the DB, so the slot must read them under its own lock rather than + # trusting this pod's in-memory router: another pod's write is invisible to that list. + fake = self._FakeDb( + [], + tuning_rows=[ + { + "model_id": row_id, + "model_name": f"router-{row_id}", + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": configs[name], + }, + } + for row_id, name in live_rows.items() + ], + ) + with ( + patch("litellm.proxy.proxy_server._license_check.auto_router_capability_limit", lambda: limit), # test-quality-ok: the guard reads the proxy license singleton with no injection seam + patch("litellm.proxy.proxy_server.llm_router", None), # test-quality-ok: the guard reads the proxy router global with no injection seam + patch("litellm.proxy.proxy_server.heuristic_v1_tuning_baselines", baselines), # test-quality-ok: baselines are a startup-loaded proxy global with no injection seam + ): + if expected == "refused": + with pytest.raises(HTTPException) as exc_info: + async with _auto_router_capability_slot(fake, effective_params=effective_params, model_id=candidate_id): + pass + assert exc_info.value.status_code == 403 + assert "changed heuristic scorer settings or tier models" in str(exc_info.value.detail) + assert "'auto_router' feature lifts the limit" in str(exc_info.value.detail) + return + async with _auto_router_capability_slot(fake, effective_params=effective_params, model_id=candidate_id) as table: + assert hasattr(table, "create") + + @pytest.mark.asyncio + async def test_add_new_model_refuses_a_second_tuned_heuristic_v1_router_without_a_model_id(self) -> None: + """A create request carries no model_info at all, yet the quota still judges it: Deployment mints the + row id before the slot is entered, so a second tuned router is refused before its DB write.""" + from litellm.proxy._types import ProxyException + from litellm.proxy.management_endpoints.model_management_endpoints import add_new_model + from litellm.router_utils.auto_router_tuning_baseline import snapshot_tuning_baselines + + admin = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) + baselines = snapshot_tuning_baselines([self._db_router_row("a", self._TUNED_A)]) + fake = self._FakeDb( + [], + tuning_rows=[ + { + "model_id": "a", + "model_name": "router-a", + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": self._TUNED_A_EDITED, + }, + } + ], + ) + with ( + patch("litellm.proxy.proxy_server.prisma_client", fake), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch("litellm.proxy.proxy_server.llm_router", None), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch("litellm.proxy.proxy_server._license_check.auto_router_capability_limit", lambda: 1), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch("litellm.proxy.proxy_server.heuristic_v1_tuning_baselines", baselines), # test-quality-ok: baselines are a startup-loaded proxy global with no injection seam + patch( # test-quality-ok: prior auth check needs a live DB; only the tuning quota is under test + "litellm.proxy.management_endpoints.model_management_endpoints.ModelManagementAuthChecks.can_user_make_model_call", + new=AsyncMock(return_value=None), + ), + patch( # test-quality-ok: params are encrypted before the slot is entered; no master key in this test + "litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper", + lambda value, new_encryption_key=None: value, + ), + ): + with pytest.raises(ProxyException) as exc_info: + await add_new_model( + model_params=Deployment( + model_name="second-tuned", + litellm_params=LiteLLM_Params( + model="auto_router/complexity_router", complexity_router_config=self._TUNED_B + ), + ), + user_api_key_dict=admin, + ) + assert exc_info.value.code == "403" + assert "changed heuristic scorer settings or tier models" in str(exc_info.value.message) + fake.tx_obj.litellm_proxymodeltable.create.assert_not_awaited() + fake.litellm_proxymodeltable.create.assert_not_awaited() + + @pytest.mark.asyncio + async def test_slot_skips_tuning_quota_when_no_baseline_is_loaded(self) -> None: + """No baseline (DB-less proxy, or the startup read failed) means the gate cannot judge, so it does not.""" + from litellm.proxy.management_endpoints.model_management_endpoints import _auto_router_capability_slot + + fake = self._FakeDb([]) + with ( + patch("litellm.proxy.proxy_server._license_check.auto_router_capability_limit", lambda: 1), # test-quality-ok: the guard reads the proxy license singleton with no injection seam + patch("litellm.proxy.proxy_server.llm_router", None), # test-quality-ok: the guard reads the proxy router global with no injection seam + patch("litellm.proxy.proxy_server.heuristic_v1_tuning_baselines", None), # test-quality-ok: baselines are a startup-loaded proxy global with no injection seam + ): + async with _auto_router_capability_slot( + fake, + effective_params={"model": "auto_router/complexity_router", "complexity_router_config": self._TUNED_B}, + model_id="c", + ) as table: + assert hasattr(table, "create") + @pytest.mark.asyncio async def test_team_model_bookkeeping_runs_after_the_slot_is_released(self) -> None: """team_model_add needs a second pool connection, so it must run only after the slot transaction diff --git a/tests/test_litellm/proxy/proxy_server/test_lifecycle.py b/tests/test_litellm/proxy/proxy_server/test_lifecycle.py index a06e7142122..a518c9289ae 100644 --- a/tests/test_litellm/proxy/proxy_server/test_lifecycle.py +++ b/tests/test_litellm/proxy/proxy_server/test_lifecycle.py @@ -61,6 +61,7 @@ def test_cleanup_router_config_variables_resets_globals(monkeypatch): monkeypatch.setattr(ps, "user_custom_auth", lambda x: x, raising=False) monkeypatch.setattr(ps, "health_check_interval", 42, raising=False) monkeypatch.setattr(ps, "prisma_client", MagicMock(), raising=False) + monkeypatch.setattr(ps, "heuristic_v1_tuning_baselines", {"router": "baseline"}, raising=False) cleanup_router_config_variables() @@ -70,6 +71,7 @@ def test_cleanup_router_config_variables_resets_globals(monkeypatch): "user_custom_auth": ps.user_custom_auth, "health_check_interval": ps.health_check_interval, "prisma_client": ps.prisma_client, + "heuristic_v1_tuning_baselines": ps.heuristic_v1_tuning_baselines, } assert normalize(observed) == { "master_key": None, @@ -77,6 +79,7 @@ def test_cleanup_router_config_variables_resets_globals(monkeypatch): "user_custom_auth": None, "health_check_interval": None, "prisma_client": None, + "heuristic_v1_tuning_baselines": None, } @@ -818,6 +821,21 @@ def test_proxy_startup_event_warns_for_global_budget_without_database(): ) +@pytest.mark.asyncio +async def test_tuning_baseline_waits_for_a_complete_db_model_census(monkeypatch): + prisma_client = MagicMock() + monkeypatch.setattr(ps.proxy_config, "_get_models_from_db", AsyncMock(return_value=None)) + + result = await ProxyStartupEvent.enforce_heuristic_v1_tuning_baseline( + prisma_client=prisma_client, + llm_router=None, + limit=1, + ) + + assert result is None + prisma_client.db.litellm_config.find_unique.assert_not_called() + + # --------------------------------------------------------------------------- # _initialize_slack_alerting_jobs — spend-report pod locking (issue #14809) # --------------------------------------------------------------------------- diff --git a/tests/test_litellm/router_utils/test_auto_router_tuning_baseline.py b/tests/test_litellm/router_utils/test_auto_router_tuning_baseline.py new file mode 100644 index 00000000000..c3395ab641f --- /dev/null +++ b/tests/test_litellm/router_utils/test_auto_router_tuning_baseline.py @@ -0,0 +1,185 @@ +"""Behavior pins for the baseline-relative heuristic-v1 tuning gate.""" + +from __future__ import annotations + +from collections.abc import Mapping + +import pytest + +from litellm.router_utils.auto_router_tuning_baseline import ( + DEFAULT_TUNING_FINGERPRINT, + HEURISTIC_V1_TUNING_FIELDS, + heuristic_v1_router_fingerprint, + mutable_tuned_identities, + router_identity, + snapshot_tuning_baselines, + tuning_fingerprint, + tuning_limit_violation, + tuning_quota_violation, +) + +_TIERS = {"SIMPLE": "cheap", "MEDIUM": "mid", "COMPLEX": "strong", "REASONING": "top"} +_ALT_TIERS = {**_TIERS, "COMPLEX": "other-strong"} + + +def _router( + name: str, + config: Mapping[str, object] | None, + *, + tags: list[str] | None = None, + db_id: str | None = None, + model: str = "auto_router/complexity_router", +) -> dict[str, object]: + litellm_params: dict[str, object] = {"model": model} + if config is not None: + litellm_params["complexity_router_config"] = dict(config) + if tags is not None: + litellm_params["tags"] = tags + row: dict[str, object] = {"model_name": name, "litellm_params": litellm_params} + if db_id is not None: + row["model_info"] = {"id": db_id, "db_model": True} + return row + + +class TestTuningFingerprint: + def test_normalized_spellings_share_one_fingerprint(self) -> None: + canonical = tuning_fingerprint({"tiers": _TIERS, "dimension_weights": {"codePresence": 0.3}}) + assert canonical == tuning_fingerprint({"dimension_weights": {"codePresence": 0.3}, "tiers": _TIERS}) + assert tuning_fingerprint({"tiers": {"SIMPLE": {"model_name": "x"}}}) == tuning_fingerprint( + {"tiers": {"SIMPLE": "x"}} + ) + + @pytest.mark.parametrize("field", sorted(set(HEURISTIC_V1_TUNING_FIELDS) - {"tier_model_configs"})) + def test_every_tuning_field_changes_the_fingerprint(self, field: str) -> None: + samples: dict[str, object] = { + "tiers": _ALT_TIERS, + "classifier_type": "heuristic_first", + "tier_boundaries": {"simple_medium": 0.2, "medium_complex": 0.4, "complex_reasoning": 0.7}, + "reasoning_override_min_score": 0.05, + "token_thresholds": {"simple": 20, "complex": 500}, + "dimension_weights": {"codePresence": 0.9}, + "code_keywords": ["orionflow"], + "reasoning_keywords": ["deduce"], + "technical_keywords": ["ledgerkit"], + "custom_technical_keywords": ["acmeflow"], + "simple_keywords": ["hey"], + "escalation_keywords": ["ESCALATE"], + "keyword_tier_rules": [{"keywords": ["urgent"], "tier": "COMPLEX"}], + } + config: dict[str, object] = {field: samples[field]} + if field == "classifier_type": + config["heuristic_first_max_tier"] = "MEDIUM" + config["classifier_llm_config"] = {"model": "judge"} + assert tuning_fingerprint(config) != DEFAULT_TUNING_FINGERPRINT + + def test_tier_model_overrides_change_the_fingerprint(self) -> None: + plain = tuning_fingerprint({"tiers": {"SIMPLE": "x"}}) + with_override = tuning_fingerprint( + {"tiers": {"SIMPLE": {"model_name": "x", "litellm_params": {"temperature": 0.1}}}} + ) + assert plain != with_override + + def test_non_tuning_fields_do_not_change_the_fingerprint(self) -> None: + assert tuning_fingerprint({"return_raw_model_name": True, "session_affinity": True}) == DEFAULT_TUNING_FINGERPRINT + + def test_invalid_config_has_no_fingerprint(self) -> None: + assert tuning_fingerprint({"tier_boundaries": "not-a-mapping"}) is None + + +class TestRouterIdentity: + def test_db_rows_key_on_model_id_and_yaml_rows_on_name_and_tags(self) -> None: + db_row = _router("renamed", {"tiers": _TIERS}, db_id="row-1") + assert router_identity(db_row) == router_identity(_router("other-name", {"tiers": _TIERS}, db_id="row-1")) + assert router_identity(_router("a", {"tiers": _TIERS}, tags=["x", "y"])) == router_identity( + _router("a", {"tiers": _TIERS}, tags=["y", "x"]) + ) + assert router_identity(_router("a", {"tiers": _TIERS}, tags=["x"])) != router_identity( + _router("a", {"tiers": _TIERS}) + ) + assert router_identity({"litellm_params": {"model": "auto_router/complexity_router"}}) is None + + +class TestHeuristicV1Scope: + @pytest.mark.parametrize( + "config,in_scope", + [ + ({"tiers": _TIERS}, True), + ({"classifier_type": "heuristic", "tiers": _TIERS}, True), + ( + { + "classifier_type": "heuristic_first", + "heuristic_first_max_tier": "MEDIUM", + "classifier_llm_config": {"model": "judge"}, + "tiers": _TIERS, + }, + True, + ), + ( + { + "classifier_type": "hybrid", + "hybrid_boundary_margin": 0.05, + "classifier_llm_config": {"model": "judge"}, + "tiers": _TIERS, + }, + True, + ), + ({"classifier_type": "heuristic_v2", "tiers": _TIERS}, False), + ({"classifier_type": "llm", "classifier_llm_config": {"model": "judge"}, "tiers": _TIERS}, False), + ], + ) + def test_only_v1_scoring_classifiers_are_fingerprinted(self, config: Mapping[str, object], in_scope: bool) -> None: + assert (heuristic_v1_router_fingerprint(_router("r", config)) is not None) is in_scope + + def test_plain_deployments_are_ignored(self) -> None: + assert heuristic_v1_router_fingerprint(_router("gpt", None, model="openai/gpt-4o")) is None + + +class TestQuota: + def test_snapshot_records_every_v1_router_even_at_defaults(self) -> None: + baselines = snapshot_tuning_baselines([_router("a", {"tiers": _TIERS}), _router("b", {})]) + assert set(baselines) == {router_identity(_router("a", {})), router_identity(_router("b", {}))} + assert baselines[router_identity(_router("b", {}))] == DEFAULT_TUNING_FINGERPRINT + + def test_unchanged_snapshot_is_never_mutable(self) -> None: + rows = [_router("a", {"tiers": _TIERS}), _router("b", {"tiers": _ALT_TIERS})] + baselines = snapshot_tuning_baselines(rows) + assert mutable_tuned_identities(rows, baselines) == frozenset() + + def test_router_added_after_snapshot_is_mutable_only_when_tuned(self) -> None: + baselines = snapshot_tuning_baselines([_router("a", {"tiers": _TIERS})]) + assert mutable_tuned_identities([_router("new", {})], baselines) == frozenset() + assert mutable_tuned_identities([_router("new", {"tiers": _TIERS})], baselines) == { + router_identity(_router("new", {})) + } + + def test_quota_matrix(self) -> None: + legacy_a = _router("a", {"tiers": _TIERS}) + legacy_b = _router("b", {"tiers": _ALT_TIERS}) + baselines = snapshot_tuning_baselines([legacy_a, legacy_b]) + edited_a = _router("a", {"tiers": _TIERS, "dimension_weights": {"codePresence": 0.9}}) + edited_b = _router("b", {"tiers": _TIERS}) + new_c = _router("c", {"tiers": _TIERS}) + + assert tuning_quota_violation(candidate=edited_a, others=[legacy_b], baselines=baselines, limit=1) is None + assert tuning_quota_violation(candidate=edited_a, others=[edited_a, legacy_b], baselines=baselines, limit=1) is None + assert tuning_quota_violation(candidate=legacy_a, others=[edited_b], baselines=baselines, limit=1) is None + assert tuning_quota_violation(candidate=edited_b, others=[edited_a], baselines=baselines, limit=1) is not None + assert tuning_quota_violation(candidate=new_c, others=[edited_a], baselines=baselines, limit=1) is not None + assert tuning_quota_violation(candidate=new_c, others=[edited_a], baselines=baselines, limit=None) is None + assert tuning_quota_violation(candidate=legacy_a, others=[edited_a, edited_b], baselines=baselines, limit=1) is None + + def test_reverting_to_baseline_frees_the_quota(self) -> None: + legacy_a = _router("a", {"tiers": _TIERS}) + legacy_b = _router("b", {"tiers": _ALT_TIERS}) + baselines = snapshot_tuning_baselines([legacy_a, legacy_b]) + edited_b = _router("b", {"tiers": _TIERS}) + assert tuning_quota_violation(candidate=edited_b, others=[legacy_a], baselines=baselines, limit=1) is None + assert tuning_quota_violation(candidate=edited_b, others=[legacy_a, edited_b], baselines=baselines, limit=1) is None + + def test_violation_message_names_the_limit_and_remedy(self) -> None: + message = tuning_limit_violation(held=2, limit=1) + assert message is not None + assert "At most 1 auto-router(s)" in message + assert "revert the other changed router to its baseline" in message + assert tuning_limit_violation(held=1, limit=1) is None + assert tuning_limit_violation(held=5, limit=None) is None