From df73c623b231b68d690f349a7bb70a05b4c82333 Mon Sep 17 00:00:00 2001 From: tin-berri Date: Thu, 3 Sep 2026 13:39:58 -0700 Subject: [PATCH] feat(router): limit heuristic_v2 auto-routers to one without the auto_router license feature (#39468) Without the auto_router feature in the signed enterprise license a proxy may hold one complexity router with classifier_type heuristic_v2 across config.yaml and the DB; with it the limit is lifted. The ceiling is derived once from LicenseCheck and handed to the Router, which refuses the extra router at registration. config.yaml over the limit refuses to start, and /model/new, /model/update and PATCH /model/{id}/update refuse the write with a 403 before touching the DB. Expiry follows the existing max_users/max_teams pattern: judged when the license is verified, not on every call, and a verify that rejects the license (expired or unreadable) leaves no signed payload behind. The rollback after a failed upsert re-admits state that was already serving, so it is exempt from the ceiling: an edit that fails, including one refused by a ceiling that has since tightened, leaves the router serving its previous configuration. A write that leaves a row on heuristic_v2 under a limited license runs in one transaction that takes a Postgres advisory lock before counting the DB rows plus this proxy's config.yaml routers, so concurrent writes on any pod cannot both claim the sole slot and no surplus row is ever persisted. Only the row insert runs under that lock: the team model bookkeeping, which needs a second pool connection, runs after the transaction has committed. PATCH /model/{id}/update follows the same order as create: the row is written through the slot first and the team's model list is updated only afterwards, so a refused write leaves the team as it was. The slot transaction bypasses the repository's publish-on-write, so it publishes the config change once after commit, as delete_team_models does. --- litellm/constants.py | 1 + litellm/proxy/auth/litellm_license.py | 23 +- .../model_management_endpoints.py | 198 +++++--- litellm/proxy/proxy_server.py | 20 +- litellm/router.py | 46 +- .../router_utils/auto_router_model_naming.py | 34 +- litellm/types/router.py | 12 + .../proxy/auth/test_litellm_license.py | 69 +++ .../test_model_management_endpoints.py | 421 ++++++++++++++++-- .../test_ptu_model_settings.py | 6 + .../proxy/proxy_server/test_proxy_config.py | 115 +++++ .../router_strategy/test_complexity_router.py | 160 +++++++ .../test_auto_router_model_naming.py | 57 +++ 13 files changed, 1067 insertions(+), 95 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 1063b6ddeeb..f5acadc32ab 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -40,6 +40,7 @@ ROUTER_SETTINGS_MANAGED_OUTSIDE_CONFIG: Final[frozenset[str]] = frozenset( "router_general_settings", "ignore_invalid_deployments", "fallback_access_check", + "heuristic_v2_router_limit", } ) DEFAULT_BATCH_SIZE: Final = int(os.getenv("DEFAULT_BATCH_SIZE", 512)) diff --git a/litellm/proxy/auth/litellm_license.py b/litellm/proxy/auth/litellm_license.py index 677f1a0fdda..55bb1e3925a 100644 --- a/litellm/proxy/auth/litellm_license.py +++ b/litellm/proxy/auth/litellm_license.py @@ -16,6 +16,10 @@ if TYPE_CHECKING: from litellm.proxy._types import EnterpriseLicenseData +AUTO_ROUTER_LICENSE_FEATURE: Final = "auto_router" +HEURISTIC_V2_LICENSE_REMEDY: Final = "A LiteLLM license with the 'auto_router' feature lifts the limit." + + class LicenseCheck: """ - Check if license in env @@ -149,6 +153,19 @@ class LicenseCheck: return False return team_count > _max_teams_in_license + def heuristic_v2_router_limit(self) -> int | None: + """ + How many heuristic_v2 auto-routers this proxy may hold: unlimited (None) only when the + signed license lists the auto_router feature, otherwise one. A license verified through + the API carries no feature list, so it does not lift the limit either. + """ + if self.airgapped_license_data is None: + return 1 + allowed_features: Final = self.airgapped_license_data.get("allowed_features") + if isinstance(allowed_features, list) and AUTO_ROUTER_LICENSE_FEATURE in allowed_features: + return None + return 1 + def verify_license_without_api_request(self, public_key, license_key): try: from cryptography.hazmat.primitives import hashes @@ -179,19 +196,21 @@ class LicenseCheck: # Decode and parse the data license_data: Final = json.loads(message.decode()) - self.airgapped_license_data = EnterpriseLicenseData(**license_data) - # debug information provided in license data verbose_proxy_logger.debug("License data: %s", license_data) # Check expiration date expiration_date: Final = datetime.strptime(license_data["expiration_date"], "%Y-%m-%d") if expiration_date < datetime.now(): + self.airgapped_license_data = None return False, "License has expired" + self.airgapped_license_data = EnterpriseLicenseData(**license_data) + return True except Exception as e: + self.airgapped_license_data = None verbose_proxy_logger.debug( "litellm.proxy.auth.litellm_license.py::verify_license_without_api_request - Unable to verify License locally. - %s", e, diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index 613e726f89d..82ee33cbc39 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -13,10 +13,11 @@ model/{model_id}/update - PATCH endpoint for model update. import asyncio import datetime import json -from collections.abc import Awaitable, Mapping, Sequence +from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping, Sequence +from contextlib import AbstractAsyncContextManager, asynccontextmanager from json import JSONDecodeError from types import MappingProxyType -from typing import TYPE_CHECKING, Annotated, Final, Literal, Protocol, cast +from typing import TYPE_CHECKING, Annotated, Final, Literal, Protocol, TypeVar, cast from fastapi import APIRouter, Depends, Header, HTTPException, Request, status from pydantic import BaseModel, ConfigDict, Field, ValidationError, field_validator @@ -49,6 +50,7 @@ from litellm.proxy._types import ( TeamModelDeleteRequest, UserAPIKeyAuth, ) +from litellm.proxy.auth.litellm_license import HEURISTIC_V2_LICENSE_REMEDY from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.config_sync_pubsub import ( coordination_redis_cache, @@ -96,6 +98,9 @@ from litellm.router_strategy.complexity_router import ( from litellm.router_utils.auto_router_model_naming import ( STRATEGY_ROUTER_PARAM_FIELDS, carries_complexity_router_settings, + count_heuristic_v2_routers, + heuristic_v2_limit_violation, + uses_heuristic_v2_classifier, validate_complexity_router_config_placement, validate_complexity_router_config_write, validate_strategy_router_model_write, @@ -153,6 +158,8 @@ class _ProxyModelTable(Protocol): def find_many(self, *, where: Mapping[str, object]) -> Awaitable[Sequence[_ProxyModelRow]]: ... + def create(self, *, data: Mapping[str, object]) -> Awaitable[_ProxyModelRow]: ... + def update( self, *, where: Mapping[str, object], data: Mapping[str, object] ) -> Awaitable[_ProxyModelRow | None]: ... @@ -166,6 +173,9 @@ class _TxModelTables(Protocol): litellm_proxymodeltable: _ProxyModelTable +_RowT = TypeVar("_RowT") + + class _ExistingModelRow(Protocol): @property def litellm_params(self) -> Mapping[str, object]: ... @@ -269,6 +279,66 @@ def _raise_on_strategy_router_write_violation( ) +HEURISTIC_V2_SLOT_LOCK_KEY: Final = 5_872_301 +_HEURISTIC_V2_LOCK_SQL: Final = "SELECT 1 AS locked FROM pg_advisory_xact_lock($1)" +_HEURISTIC_V2_DB_ROWS_SQL: Final = """ +SELECT count(*)::int AS held FROM "LiteLLM_ProxyModelTable" +WHERE model_id <> $1 + AND (CASE jsonb_typeof(litellm_params) WHEN 'string' THEN (litellm_params #>> '{}')::jsonb ELSE litellm_params END) + -> 'complexity_router_config' ->> 'classifier_type' = 'heuristic_v2' +""" + + +def _effective_complexity_router_config( + incoming_params: GenericLiteLLMParams | None, existing_params: GenericLiteLLMParams | None +) -> object: + """The complexity config a write leaves on the row: the incoming one when the write carries it, else the stored one.""" + incoming: Final = None if incoming_params is None else incoming_params.complexity_router_config + if incoming is not None or existing_params is None: + return incoming + return existing_params.complexity_router_config + + +@asynccontextmanager +async def _heuristic_v2_slot( + prisma_client: PrismaClient, *, effective_config: object, model_id: str | None +) -> AsyncGenerator[_ProxyModelTable, None]: + """Hand out the model table to write through while the row's claim on a heuristic_v2 slot is settled. + + A write that leaves the row on classifier_type heuristic_v2 under a limited license runs + inside one transaction that takes an advisory lock in its own statement before counting + (a statement's snapshot predates anything it locks), so pods cannot both pass the count: + the DB rows (any pod, either JSON shape) plus this proxy's config.yaml routers are judged + against the license limit and the write is refused with a 403 before it happens. The row + being edited keeps its own slot through ``model_id``. Every other write, and every write on + an unlimited license, goes through the repository table with no lock. Only the row write + itself may run inside: anything that needs a second connection (the team model bookkeeping) + 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. + """ + from litellm.proxy.proxy_server import _license_check, llm_router + + limit: Final = _license_check.heuristic_v2_router_limit() + if limit is None or not uses_heuristic_v2_classifier(effective_config): + 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(_HEURISTIC_V2_LOCK_SQL, HEURISTIC_V2_SLOT_LOCK_KEY) + rows: Sequence[Mapping[str, object]] = await tx_ctx.query_raw(_HEURISTIC_V2_DB_ROWS_SQL, model_id or "") + db_held: Final = rows[0].get("held") if rows else 0 + config_rows: Final = () if llm_router is None else tuple(llm_router.config_deployments()) + held: Final = (db_held if isinstance(db_held, int) else 0) + count_heuristic_v2_routers(config_rows) + violation: Final = heuristic_v2_limit_violation(held=held + 1, limit=limit) + if violation is not None: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, detail=f"{violation} {HEURISTIC_V2_LICENSE_REMEDY}" + ) + yield tables.litellm_proxymodeltable + await publish_config_change(redis_cache=coordination_redis_cache(), object_type="litellm_proxymodeltable") + + ENFORCE_RPM_TPM_ON_MODEL_ADD_SETTING: Final = "enforce_rpm_tpm_on_model_add" _REQUIRED_RATE_LIMIT_FIELDS: Final = ("rpm", "tpm") @@ -720,22 +790,29 @@ async def patch_model( ) requested_model_name: Final = patch_data.model_name + stored_model_name: str | None = None + + async def write_row(update_data: PrismaCompatibleUpdateDBModel) -> _ProxyModelRow | None: + nonlocal stored_model_name + stored_model_name = update_data.get("model_name") + update_data["updated_by"] = user_api_key_dict.user_id or litellm_proxy_admin_name + update_data["updated_at"] = cast(str, get_utc_datetime()) + async with _heuristic_v2_slot( + prisma_client, + effective_config=_effective_complexity_router_config( + patch_data.litellm_params, db_model.litellm_params + ), + model_id=model_id, + ) as table: + return await table.update(where={"model_id": model_id}, data=update_data) + # Handle team model updates with proper alias management - update_data: Final = await _update_team_model_in_db( + updated_model: Final = await _update_team_model_in_db( db_model=db_model, patch_data=patch_data, user_api_key_dict=user_api_key_dict, prisma_client=prisma_client, - ) - - # Add metadata about update - update_data["updated_by"] = user_api_key_dict.user_id or litellm_proxy_admin_name - update_data["updated_at"] = cast(str, get_utc_datetime()) - - # Perform partial update - updated_model: Final = await _proxy_model_table(prisma_client).update( - where={"model_id": model_id}, - data=update_data, + write_row=write_row, ) if updated_model is None: @@ -746,7 +823,6 @@ async def patch_model( param=None, ) - stored_model_name: Final = update_data.get("model_name") if ( stored_model_name is not None and stored_model_name == requested_model_name @@ -980,7 +1056,8 @@ async def _add_model_to_db( prisma_client: PrismaClient, new_encryption_key: str | None = None, should_create_model_in_db: bool = True, -) -> "prisma_models.LiteLLM_ProxyModelTable | LiteLLM_ProxyModelTable | None": + slot: AbstractAsyncContextManager[_ProxyModelTable] | None = None, +) -> "_ProxyModelRow | LiteLLM_ProxyModelTable": # encrypt litellm params # _litellm_params_dict: Final = model_params.litellm_params.dict(exclude_none=True) _original_litellm_model_name: Final = model_params.litellm_params.model @@ -998,18 +1075,20 @@ async def _add_model_to_db( if model_params.model_info.id is not None: _data["model_id"] = model_params.model_info.id _create_data: Final = cast("Mapping[str, object]", _data) # cast-ok: str-keyed json payload built just above - if should_create_model_in_db: - model_response = await ModelRepository(prisma_client).table.create(data=_create_data) - else: - model_response = LiteLLM_ProxyModelTable(**_data) - return model_response + if not should_create_model_in_db: + return LiteLLM_ProxyModelTable(**_data) + if slot is None: + return await _proxy_model_table(prisma_client).create(data=_create_data) + async with slot as table: + return await table.create(data=_create_data) async def _add_team_model_to_db( model_params: Deployment, user_api_key_dict: UserAPIKeyAuth, prisma_client: PrismaClient, -) -> "prisma_models.LiteLLM_ProxyModelTable | LiteLLM_ProxyModelTable | None": + slot: AbstractAsyncContextManager[_ProxyModelTable] | None = None, +) -> "_ProxyModelRow | LiteLLM_ProxyModelTable": """ If 'team_id' is provided, @@ -1040,6 +1119,7 @@ async def _add_team_model_to_db( model_params=model_params, user_api_key_dict=user_api_key_dict, prisma_client=prisma_client, + slot=slot, ) if original_model_name: @@ -1060,7 +1140,8 @@ async def _update_team_model_in_db( patch_data: updateDeployment, user_api_key_dict: UserAPIKeyAuth, prisma_client: PrismaClient, -) -> PrismaCompatibleUpdateDBModel: + write_row: Callable[[PrismaCompatibleUpdateDBModel], Awaitable[_RowT]], +) -> _RowT: """ Handle team model updates with proper alias management. @@ -1068,6 +1149,9 @@ async def _update_team_model_in_db( - Creates unique internal model_name and team alias - Adds model to team object - Preserves team_public_model_name for external reference + + The row is written through ``write_row`` before the team's model list is touched, so a + refused or failed write leaves the team as it was (the create path orders itself the same way). """ # Validate team_id if present in patch_data from litellm.proxy.proxy_server import premium_user @@ -1079,9 +1163,7 @@ async def _update_team_model_in_db( premium_user=premium_user, ) - # Validated before any write, beside the premium check the create path already runs - # here. The team ACL is updated below and autocommits, so a validator that raises - # further down would leave the team mutated and the deployment row never written. + # Validated before the row write, beside the premium check the create path already runs here. # # The merged view is what gets stored, so that is what has to satisfy the invariants. # Validating the patch alone rejected a partial edit of an already valid deployment: @@ -1101,7 +1183,7 @@ async def _update_team_model_in_db( # No team_id in patch, proceed with standard update if patch_team_id is None: - return update_db_model(db_model=db_model, updated_patch=patch_data) + return await write_row(update_db_model(db_model=db_model, updated_patch=patch_data)) # Determine public model name public_model_name: Final = _get_public_model_name( @@ -1120,11 +1202,14 @@ async def _update_team_model_in_db( db_team_id: Final = db_model.model_info.team_id if db_model.model_info else None is_new_team_assignment: Final = db_team_id != patch_team_id + # Team rows keep their internal UUID-based model_name; the public name lives in model_info + patch_data.model_name = f"model_name_{patch_team_id}_{uuid.uuid4()}" if is_new_team_assignment else None + row: Final = await write_row(update_db_model(db_model=db_model, updated_patch=patch_data)) + if is_new_team_assignment: await _setup_new_team_model_assignment( team_id=patch_team_id, public_model_name=public_model_name, - patch_data=patch_data, user_api_key_dict=user_api_key_dict, ) else: @@ -1132,12 +1217,11 @@ async def _update_team_model_in_db( team_id=patch_team_id, public_model_name=public_model_name, db_model=db_model, - patch_data=patch_data, user_api_key_dict=user_api_key_dict, prisma_client=prisma_client, ) - return update_db_model(db_model=db_model, updated_patch=patch_data) + return row def _get_public_model_name( @@ -1189,13 +1273,9 @@ def _get_public_model_name( async def _setup_new_team_model_assignment( team_id: str, public_model_name: str, - patch_data: updateDeployment, user_api_key_dict: UserAPIKeyAuth, ) -> None: - """Set up a new team model with unique name and team membership.""" - unique_model_name: Final = f"model_name_{team_id}_{uuid.uuid4()}" - patch_data.model_name = unique_model_name - + """Register a newly team-assigned model's public name on the team.""" await team_model_add( data=TeamModelAddRequest( team_id=team_id, @@ -1385,7 +1465,6 @@ async def _update_existing_team_model_assignment( team_id: str, public_model_name: str, db_model: Deployment, - patch_data: updateDeployment, user_api_key_dict: UserAPIKeyAuth, prisma_client: PrismaClient | None, ) -> None: @@ -1409,9 +1488,6 @@ async def _update_existing_team_model_assignment( old_public_name: Final = db_model.model_info.team_public_model_name if db_model.model_info else None if old_public_name and public_model_name != old_public_name: - # Clear user-supplied public name from patch before any early return so the - # caller does not overwrite the internal UUID-based model_name in the DB. - patch_data.model_name = None if prisma_client is None: verbose_proxy_logger.warning( "prisma_client not initialized; skipping public name update entirely to avoid orphaned entries" @@ -1459,10 +1535,6 @@ async def _update_existing_team_model_assignment( # else: old_public_name == public_model_name (no rename needed) # No team_model_add/delete calls required; public name is already registered - # Always clear patch_data.model_name to prevent caller from overwriting - # the internal UUID-based model_name in the DB with the user-supplied public name - patch_data.model_name = None - class ModelManagementAuthChecks: """ @@ -1878,18 +1950,19 @@ async def add_new_model( reload_outcome: ReconcileOutcome = ReconcileOutcome(still_desired=None, live_after=None) try: _original_litellm_model_name: Final = model_params.model_name - if model_params.model_info.team_id is None: - model_response = await _add_model_to_db( - model_params=priced_model_params, - user_api_key_dict=user_api_key_dict, - prisma_client=prisma_client, - ) - else: - model_response = await _add_team_model_to_db( - model_params=priced_model_params, - user_api_key_dict=user_api_key_dict, - prisma_client=prisma_client, - ) + add_model: Final = ( + _add_model_to_db if model_params.model_info.team_id is None else _add_team_model_to_db + ) + model_response = await add_model( + model_params=priced_model_params, + user_api_key_dict=user_api_key_dict, + prisma_client=prisma_client, + slot=_heuristic_v2_slot( + prisma_client, + effective_config=priced_model_params.litellm_params.complexity_router_config, + model_id=priced_model_params.model_info.id, + ), + ) reload_outcome = await proxy_config.add_deployment( prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj ) @@ -1903,6 +1976,8 @@ async def add_new_model( passed_model_info=priced_model_params.model_info, ) except Exception as e: + if isinstance(e, HTTPException): + raise verbose_proxy_logger.exception("Exception in add_new_model: %s", e) else: @@ -2070,10 +2145,17 @@ async def update_model( "updated_by": user_api_key_dict.user_id or LITELLM_PROXY_ADMIN_NAME, **({} if renamed_to is None else {"model_name": renamed_to}), } - model_response: Final = await _proxy_model_table(prisma_client).update( - where={"model_id": _model_id}, - data=_data, - ) + async with _heuristic_v2_slot( + prisma_client, + effective_config=_effective_complexity_router_config( + model_params.litellm_params, deployment.litellm_params + ), + model_id=_model_id, + ) as table: + model_response: Final = await table.update( + where={"model_id": _model_id}, + data=_data, + ) if renamed_to is not None: await sync_access_groups_for_renamed_model( prisma_client=prisma_client, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 27132c90e05..96b425f2a24 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -120,6 +120,8 @@ from litellm.router_utils.add_retry_fallback_headers import ( from litellm.router_utils.auto_router_model_naming import ( STRATEGY_ROUTER_PARAM_FIELDS, carries_complexity_router_settings, + count_heuristic_v2_routers, + heuristic_v2_limit_violation, validate_complexity_router_config_placement, ) from litellm.types.utils import ( @@ -301,7 +303,7 @@ from litellm.proxy.auth.auth_utils import ( ) from litellm.proxy.auth.fallback_model_access import router_fallback_access_check from litellm.proxy.auth.handle_jwt import JWTHandler -from litellm.proxy.auth.litellm_license import LicenseCheck +from litellm.proxy.auth.litellm_license import HEURISTIC_V2_LICENSE_REMEDY, LicenseCheck from litellm.proxy.auth.model_checks import ( expand_wildcard_deployments_for_model_info, get_all_fallbacks, @@ -4316,6 +4318,19 @@ def validate_deployment_complexity_router_placement(model: Mapping[str, object]) raise ValueError(f"model {model.get('model_name', '')!r}: {violation}") +def validate_heuristic_v2_router_limit(model_list: Sequence[Mapping[str, object]], *, limit: int | None) -> None: + """ + Refuse to start when config.yaml defines more heuristic_v2 auto-routers than the license allows. + + Checked here rather than left to router registration for the same reason as the two + validators above: the proxy builds its router with `ignore_invalid_deployments=True`, so + the router's own refusal would turn the extra router into a silently missing model. + """ + violation: Final = heuristic_v2_limit_violation(held=count_heuristic_v2_routers(model_list), limit=limit) + if violation is not None: + raise ValueError(f"config.yaml model_list: {violation} {HEURISTIC_V2_LICENSE_REMEDY}") + + def pin_complexity_router_model_id(model: dict) -> None: # mutable-ok: out-param, model_info is stamped in place """ Stamps `model_info.id` from the raw litellm_params before plugin resolution swaps @@ -5721,6 +5736,7 @@ class ProxyConfig: model_list: Final = config.get("model_list", None) if model_list: router_params["model_list"] = model_list + validate_heuristic_v2_router_limit(model_list, limit=_license_check.heuristic_v2_router_limit()) print( # noqa: T201 "\033[32mLiteLLM: Proxy initialized with Config, Set models:\033[0m" ) @@ -5810,6 +5826,7 @@ class ProxyConfig: ), ignore_invalid_deployments=True, # don't raise an error if a deployment is invalid fallback_access_check=router_fallback_access_check, + heuristic_v2_router_limit=_license_check.heuristic_v2_router_limit, ) if redis_usage_cache is not None and router.cache.redis_cache is None: @@ -6270,6 +6287,7 @@ class ProxyConfig: search_tools=search_tools, ignore_invalid_deployments=True, fallback_access_check=router_fallback_access_check, + heuristic_v2_router_limit=_license_check.heuristic_v2_router_limit, ) verbose_proxy_logger.debug("updated llm_router: %s", llm_router) else: diff --git a/litellm/router.py b/litellm/router.py index cfb080a24a0..dc750941559 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -21,7 +21,7 @@ import time import traceback import weakref from collections import defaultdict -from collections.abc import AsyncGenerator, AsyncIterator, Callable, Generator, Mapping, Sequence +from collections.abc import AsyncGenerator, AsyncIterator, Callable, Generator, Iterator, Mapping, Sequence from functools import lru_cache, partial from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypeAlias, TypeVar, Union, cast @@ -117,6 +117,9 @@ from litellm.router_utils.add_retry_fallback_headers import ( from litellm.router_utils.auto_router_model_naming import ( AUTO_ROUTER_MODEL_PREFIX, classify_strategy_router_model, + count_heuristic_v2_routers, + heuristic_v2_limit_violation, + uses_heuristic_v2_classifier, ) from litellm.router_utils.batch_utils import ( _get_router_metadata_variable_name, @@ -211,6 +214,7 @@ from litellm.types.router import ( DeploymentTypedDict, FallbackAccessCheck, GuardrailTypedDict, + HeuristicV2RouterLimit, LiteLLM_Params, MockRouterTestingParams, ModelGroupInfo, @@ -683,6 +687,7 @@ class Router: background_health_check_model_groups: Sequence[str] | None = None, enable_weighted_failover: bool = False, fallback_access_check: FallbackAccessCheck | None = None, + heuristic_v2_router_limit: HeuristicV2RouterLimit | None = None, ) -> None: """ Initialize the Router class with the given parameters for caching, reliability, and routing strategy. @@ -759,6 +764,7 @@ class Router: self.set_verbose = set_verbose self.ignore_invalid_deployments = ignore_invalid_deployments + self.heuristic_v2_router_limit = heuristic_v2_router_limit self.fallback_access_check: Final = fallback_access_check self.debug_level = debug_level self.enable_pre_call_checks = enable_pre_call_checks @@ -8796,6 +8802,30 @@ class Router: """ return classify_strategy_router_model(litellm_params.model) == "complexity" + def config_deployments(self) -> Iterator[Mapping[str, object]]: + """The model_list rows that came from config.yaml rather than the DB (``model_info.db_model`` unset).""" + for deployment in self.model_list: + if not isinstance(deployment, Mapping): + continue + model_info = deployment.get("model_info") + if not (isinstance(model_info, Mapping) and model_info.get("db_model")): + yield deployment + + def heuristic_v2_router_limit_violation(self) -> str | None: + """ + Why one more heuristic_v2 router cannot join this router, or None when it can. + + Judged against every deployment currently on the model_list; an upsert pops the row being + edited first, so an edit of an existing heuristic_v2 router keeps its own slot. The limit is + resolved on every call through ``heuristic_v2_router_limit``; unset means unlimited, which + is the SDK default, and the proxy injects a resolver backed by its license. + """ + limit: Final = self.heuristic_v2_router_limit() if self.heuristic_v2_router_limit is not None else None + others: Final = count_heuristic_v2_routers( + deployment for deployment in self.model_list if isinstance(deployment, Mapping) + ) + return heuristic_v2_limit_violation(held=others + 1, limit=limit) + def init_complexity_router_deployment(self, deployment: Deployment): """ Initialize the complexity-router deployment. @@ -8813,6 +8843,10 @@ class Router: ) complexity_router_config: Final[dict | None] = deployment.litellm_params.complexity_router_config + if uses_heuristic_v2_classifier(complexity_router_config): + limit_violation: Final = self.heuristic_v2_router_limit_violation() + if limit_violation is not None: + raise ValueError(limit_violation) default_model: str | None = deployment.litellm_params.complexity_router_default_model @@ -9636,8 +9670,16 @@ class Router: raise e def _restore_deployment_after_failed_upsert(self, previous_deployment: Deployment | None, model_id: str) -> None: + """Put a deployment back the way it was before a failed upsert popped it. + + A rollback re-admits state that was already serving, so it does not go through the + heuristic_v2 ceiling a newcomer gets: with the ceiling tightened since the deployment first + registered, judging the rollback would drop a serving router over an unrelated failed edit. + """ if previous_deployment is None or self.has_model_id(model_id): return + limit_resolver: Final = self.heuristic_v2_router_limit + self.heuristic_v2_router_limit = None try: self.add_deployment(deployment=previous_deployment) verbose_router_logger.info( @@ -9652,6 +9694,8 @@ class Router: model_id, restore_error, ) + finally: + self.heuristic_v2_router_limit = limit_resolver @staticmethod def _backend_cost_map_keys(model: str, custom_llm_provider: str | None) -> tuple[str, ...]: diff --git a/litellm/router_utils/auto_router_model_naming.py b/litellm/router_utils/auto_router_model_naming.py index a8aa543d735..2efbfb5782e 100644 --- a/litellm/router_utils/auto_router_model_naming.py +++ b/litellm/router_utils/auto_router_model_naming.py @@ -10,7 +10,7 @@ the router silently dropping the deployment at load time under ``ignore_invalid_deployments``. """ -from collections.abc import Mapping, Sequence +from collections.abc import Iterable, Mapping, Sequence from dataclasses import dataclass from types import MappingProxyType from typing import Final, Literal, TypeAlias @@ -163,6 +163,38 @@ def strategy_router_dependencies( ) +def uses_heuristic_v2_classifier(complexity_router_config: object) -> bool: + """Whether this complexity config classifies with the bundled heuristic_v2 model.""" + return _mapping(complexity_router_config).get("classifier_type") == "heuristic_v2" + + +def is_heuristic_v2_router(litellm_params: Mapping[str, object]) -> bool: + """Whether this deployment is a complexity router that classifies with heuristic_v2.""" + return classify_strategy_router_model(str(litellm_params.get("model") or "")) == "complexity" and ( + uses_heuristic_v2_classifier(litellm_params.get("complexity_router_config")) + ) + + +def count_heuristic_v2_routers(deployments: Iterable[Mapping[str, object]]) -> int: + """How many of ``deployments`` (router model_list entries or config.yaml rows) are heuristic_v2 routers.""" + return sum(1 for deployment in deployments if is_heuristic_v2_router(_mapping(deployment.get("litellm_params")))) + + +def heuristic_v2_limit_violation(*, held: int, limit: int | None) -> str | None: + """Why holding ``held`` heuristic_v2 routers exceeds ``limit``, or None when it fits. + + ``limit`` None means unlimited. The message is shared by every enforcement point (config + load, model writes, router registration) and stays SDK-neutral: it names the cap and what + the caller can change; the proxy appends how its license lifts the cap. + """ + if limit is None or held <= limit: + return None + return ( + f"At most {limit} auto-router(s) with classifier_type 'heuristic_v2' can be registered but this would make " + f"{held}. Use classifier_type 'heuristic' for this router or remove an existing heuristic_v2 router." + ) + + def validate_complexity_router_config_write(complexity_router_config: Mapping[str, object] | None) -> str | None: """Reject a complexity config the router would refuse to build a deployment from. diff --git a/litellm/types/router.py b/litellm/types/router.py index 4f4df1a8d2e..7ebd50f1328 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -885,6 +885,18 @@ class FallbackAccessCheck(Protocol): async def __call__(self, *, model: str, request_kwargs: Mapping[str, object], llm_router: "Router") -> bool: ... +class HeuristicV2RouterLimit(Protocol): + """ + Resolves how many heuristic_v2 complexity routers the Router may hold right now; None means unlimited. + + The Router calls it on every registration and limit query instead of caching the answer, so the + proxy can keep the limit on its license object (re-verified on config load) rather than hand + over a snapshot. + """ + + def __call__(self) -> int | None: ... + + class LiteLLM_RouterFileObject(TypedDict, total=False): """ Tracking the litellm params hash, used for mapping the file id to the right model diff --git a/tests/test_litellm/proxy/auth/test_litellm_license.py b/tests/test_litellm/proxy/auth/test_litellm_license.py index 8da365cb587..1db53638070 100644 --- a/tests/test_litellm/proxy/auth/test_litellm_license.py +++ b/tests/test_litellm/proxy/auth/test_litellm_license.py @@ -2,6 +2,8 @@ import asyncio import json from unittest.mock import AsyncMock, MagicMock, patch +from cryptography.hazmat.primitives.asymmetric.rsa import RSAPublicKey + from litellm.proxy.auth.litellm_license import LicenseCheck @@ -30,3 +32,70 @@ def test_is_over_limit(): assert license_check.is_over_limit(101) is False assert license_check.is_over_limit(100) is False assert license_check.is_over_limit(99) is False + + +def test_heuristic_v2_router_limit() -> None: + """Only the signed license's auto_router feature lifts the one-router limit; an API-verified + license (no airgapped data) and an airgapped license without the feature keep it.""" + license_check = LicenseCheck() + license_check.airgapped_license_data = {"expiration_date": "2999-01-01", "allowed_features": ["auto_router"]} + assert license_check.heuristic_v2_router_limit() is None + + license_check.airgapped_license_data = { + "expiration_date": "2999-01-01", + "allowed_features": ["sso", "auto_router", "audit_logs"], + } + assert license_check.heuristic_v2_router_limit() is None + + license_check.airgapped_license_data = {"expiration_date": "2999-01-01", "allowed_features": ["sso"]} + assert license_check.heuristic_v2_router_limit() == 1 + + license_check.airgapped_license_data = {"expiration_date": "2999-01-01"} + assert license_check.heuristic_v2_router_limit() == 1 + + license_check.airgapped_license_data = None + assert license_check.heuristic_v2_router_limit() == 1 + + +def _signed_license(expiration_date: str) -> tuple[RSAPublicKey, str]: + import base64 + + from cryptography.hazmat.primitives import hashes + from cryptography.hazmat.primitives.asymmetric import padding, rsa + + private_key = rsa.generate_private_key(public_exponent=65537, key_size=2048) + message = json.dumps( + {"expiration_date": expiration_date, "user_id": "u", "allowed_features": ["auto_router"]} + ).encode() + signature = private_key.sign( + message, + padding.PSS(mgf=padding.MGF1(hashes.SHA256()), salt_length=padding.PSS.MAX_LENGTH), + hashes.SHA256(), + ) + return private_key.public_key(), base64.b64encode(message + b"." + signature).decode() + + +def test_expired_or_unreadable_license_grants_no_features() -> None: + """The verifier stores the signed payload only after the expiry check passes and clears it when a + later verify rejects the license, so a stale payload cannot keep lifting the heuristic_v2 limit.""" + license_check = LicenseCheck() + public_key, valid_key = _signed_license("2999-01-01") + assert license_check.verify_license_without_api_request(public_key=public_key, license_key=valid_key) is True + assert license_check.heuristic_v2_router_limit() is None + + _, expired_key = _signed_license("2000-01-01") + assert license_check.verify_license_without_api_request(public_key=public_key, license_key=expired_key) is not True + assert license_check.airgapped_license_data is None + assert license_check.heuristic_v2_router_limit() == 1 + + assert license_check.verify_license_without_api_request(public_key=public_key, license_key=valid_key) is True + assert license_check.verify_license_without_api_request(public_key=public_key, license_key="not-a-license") is not True + assert license_check.airgapped_license_data is None + + +def test_valid_signed_license_with_auto_router_lifts_the_limit() -> None: + license_check = LicenseCheck() + public_key, license_key = _signed_license("2999-01-01") + + assert license_check.verify_license_without_api_request(public_key=public_key, license_key=license_key) is True + assert license_check.heuristic_v2_router_limit() is None 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 5fa59a85c9d..c69f8f20a13 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 @@ -28,9 +28,18 @@ from litellm.proxy.management_endpoints.model_management_endpoints import ( delete_team_models, ) from litellm.proxy.utils import PrismaClient +from litellm.router import Router from litellm.types.router import Deployment, LiteLLM_Params, updateDeployment +async def _passthrough_row(update_data): + return update_data + + +async def _write_empty_row(**kwargs): + return await kwargs["write_row"]({}) + + class MockPrismaClient: def __init__( self, @@ -1191,7 +1200,7 @@ class TestTeamModelSiblingRouting: team_id = "team_no_alias" public_name = "gpt-4.1-mini" - async def mock_add_model_to_db(model_params, user_api_key_dict, prisma_client): + async def mock_add_model_to_db(model_params, user_api_key_dict, prisma_client, slot=None): return MagicMock(model_id=str(uuid.uuid4())) mock_team_model_add = AsyncMock() @@ -1372,7 +1381,8 @@ class TestTeamModelUpdate: db_model=db_model, patch_data=patch_data, user_api_key_dict=user_api_key_dict, - prisma_client=prisma_client, # type: ignore + prisma_client=prisma_client, # type: ignore, + write_row=_passthrough_row, ) assert result.get("model_name", "").startswith("model_name_test_team_123_") @@ -1435,7 +1445,6 @@ class TestTeamModelUpdate: team_id="team_123", public_model_name="new-public-name", db_model=db_model, - patch_data=patch_data, user_api_key_dict=user_api_key_dict, prisma_client=prisma_client, # type: ignore ) @@ -1481,7 +1490,6 @@ class TestTeamModelUpdate: team_id="team_123", public_model_name="new-public-name", db_model=db_model, - patch_data=patch_data, user_api_key_dict=user_api_key_dict, prisma_client=None, ) @@ -1490,39 +1498,72 @@ class TestTeamModelUpdate: mock_delete.assert_not_called() @pytest.mark.asyncio - async def test_rename_with_prisma_none_clears_patch_model_name(self): - """Rename path must clear patch_data.model_name even when prisma is unavailable (P1).""" + async def test_a_refused_row_write_leaves_the_team_untouched(self): + """The team's model list autocommits, so it is written only after the row write succeeded: a + refused write (the heuristic_v2 slot 403, a DB error) must not leave the team listing a name + whose row never changed.""" + from fastapi import HTTPException + from litellm.proxy.management_endpoints.model_management_endpoints import ( - _update_existing_team_model_assignment, + _update_team_model_in_db, ) from litellm.types.router import ModelInfo db_model = Deployment( - model_name="model_name_team_123_uuid1", + model_name="gpt-4o", litellm_params=LiteLLM_Params(model="azure/gpt-4o-mini"), - model_info=ModelInfo( - team_id="team_123", team_public_model_name="old-public-name" + model_info=ModelInfo(), + ) + user_api_key_dict = UserAPIKeyAuth(user_id="test_user", user_role=LitellmUserRoles.PROXY_ADMIN) + events: list[str] = [] + written: dict[str, object] = {} + + def patch_data() -> updateDeployment: + return updateDeployment(model_name="team-public", model_info=ModelInfo(team_id="team_123")) + + async def refuse_row(update_data): + events.append("row") + raise HTTPException(status_code=403, detail="slot held") + + async def accept_row(update_data): + events.append("row") + written.update(update_data) + return update_data + + async def team_add(**_): + events.append("team_model_add") + + with ( + patch( # test-quality-ok: the team auth check needs a live DB; the write order is what is under test + "litellm.proxy.management_endpoints.model_management_endpoints.ModelManagementAuthChecks.allow_team_model_action", + AsyncMock(return_value=True), ), - ) - patch_data = updateDeployment( - model_name="new-public-name", - model_info=ModelInfo(team_id="team_123"), - ) - user_api_key_dict = UserAPIKeyAuth( - user_id="test_user", - user_role=LitellmUserRoles.PROXY_ADMIN, - ) + patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: team models are premium-gated through a proxy global with no injection seam + patch( # test-quality-ok: the team list write is the collaborator whose ordering is asserted + "litellm.proxy.management_endpoints.model_management_endpoints.team_model_add", + side_effect=team_add, + ), + ): + with pytest.raises(HTTPException): + await _update_team_model_in_db( + db_model=db_model, + patch_data=patch_data(), + user_api_key_dict=user_api_key_dict, + prisma_client=MockPrismaClient(team_exists=True), # type: ignore + write_row=refuse_row, + ) + assert events == ["row"] - await _update_existing_team_model_assignment( - team_id="team_123", - public_model_name="new-public-name", - db_model=db_model, - patch_data=patch_data, - user_api_key_dict=user_api_key_dict, - prisma_client=None, - ) - - assert patch_data.model_name is None + await _update_team_model_in_db( + db_model=db_model, + patch_data=patch_data(), + user_api_key_dict=user_api_key_dict, + prisma_client=MockPrismaClient(team_exists=True), # type: ignore + write_row=accept_row, + ) + assert events == ["row", "row", "team_model_add"] + assert str(written["model_name"]).startswith("model_name_team_123_") + assert "team-public" in str(written["model_info"]) @pytest.mark.asyncio async def test_rename_handles_legacy_string_model_info(self): @@ -1574,7 +1615,6 @@ class TestTeamModelUpdate: team_id="team_123", public_model_name="new-public-name", db_model=db_model, - patch_data=patch_data, user_api_key_dict=user_api_key_dict, prisma_client=prisma_client, # type: ignore ) @@ -1614,7 +1654,8 @@ class TestTeamModelUpdate: db_model=db_model, patch_data=patch_data, user_api_key_dict=user_api_key_dict, - prisma_client=prisma_client, # type: ignore + prisma_client=prisma_client, # type: ignore, + write_row=_passthrough_row, ) assert "403" in str(exc_info.value) @@ -1900,7 +1941,8 @@ class TestTeamModelUpdate: db_model=db_model, patch_data=patch_data, user_api_key_dict=user_api_key_dict, - prisma_client=prisma_client, # type: ignore + prisma_client=prisma_client, # type: ignore, + write_row=_passthrough_row, ) # team ACL must not be touched on a no-op edit @@ -4311,6 +4353,321 @@ class TestStrategyRouterWriteValidation: is None ) + @staticmethod + def _live_router_holding_one_heuristic_v2(limit: int | None) -> Router: + return Router( + model_list=[ + {"model_name": "gpt-4o-mini", "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "k"}}, + { + "model_name": "held-v2", + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": {"classifier_type": "heuristic_v2", "tiers": {"SIMPLE": "gpt-4o-mini"}}, + }, + "model_info": {"id": "held-id"}, + }, + ], + heuristic_v2_router_limit=lambda: limit, + ) + + class _FakeTx: + """Stands in for a prisma transaction: records the raw statements and exposes the model table.""" + + def __init__(self, db_held: int) -> None: + self.db_held = db_held + self.raw_calls: list[tuple[str, tuple[object, ...]]] = [] + self.litellm_proxymodeltable = MagicMock(create=AsyncMock(), update=AsyncMock()) + + async def query_raw(self, sql: str, *args: object) -> list[dict[str, object]]: + self.raw_calls.append((sql, args)) + return [{"held": self.db_held}] if "count(*)" in sql else [] + + async def __aenter__(self) -> "TestStrategyRouterWriteValidation._FakeTx": + return self + + async def __aexit__(self, *exc: object) -> None: + return None + + class _FakeDb: + """Stands in for prisma_client: the plain client and the transaction it opens are told apart by identity.""" + + def __init__(self, db_held: int, existing_row: object = None) -> None: + self.db = self + self.tx_obj = TestStrategyRouterWriteValidation._FakeTx(db_held) + self.litellm_proxymodeltable = MagicMock( + create=AsyncMock(), update=AsyncMock(), find_unique=AsyncMock(return_value=existing_row) + ) + + def tx(self) -> "TestStrategyRouterWriteValidation._FakeTx": + return self.tx_obj + + _V2 = {"classifier_type": "heuristic_v2", "tiers": {"SIMPLE": "gpt-4o-mini"}} + _V1 = {"classifier_type": "heuristic", "tiers": {"SIMPLE": "gpt-4o-mini"}} + + @pytest.mark.parametrize( + "incoming,existing,expected", + [ + (_V2, None, _V2), + (_V2, _V1, _V2), + (None, _V1, _V1), + (None, None, None), + ("no-config", _V2, _V2), + ], + ) + def test_effective_complexity_router_config( + self, incoming: object, existing: object, expected: object + ) -> None: + """A write is judged on the config it leaves on the row: the incoming one when it carries one, else the stored one.""" + from litellm.proxy.management_endpoints.model_management_endpoints import ( + _effective_complexity_router_config, + ) + from litellm.types.router import updateLiteLLMParams + + incoming_params = None if incoming is None else updateLiteLLMParams( + complexity_router_config=None if incoming == "no-config" else incoming + ) + existing_params = None if existing is None else updateLiteLLMParams(complexity_router_config=existing) + assert _effective_complexity_router_config(incoming_params, existing_params) == expected + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "limit,effective_config,db_held,config_holds_one,model_id,expected", + [ + (1, _V2, 1, False, None, "refused"), + (1, _V2, 0, True, None, "refused"), + (1, _V2, 0, False, None, "reserved"), + (1, _V2, 0, False, "held-id", "reserved"), + (2, _V2, 1, False, None, "reserved"), + (1, _V1, 5, True, None, "plain"), + (1, None, 5, True, None, "plain"), + (None, _V2, 5, True, None, "plain"), + ], + ) + async def test_heuristic_v2_slot_matrix( + self, + limit: int | None, + effective_config: object, + db_held: int, + config_holds_one: bool, + model_id: str | None, + expected: str, + ) -> None: + """The slot is claimed inside a locked transaction only for a heuristic_v2 write under a limit; the DB rows + (other pods included) plus config.yaml routers decide, the row being edited is excluded through the SQL + parameter, and every other write runs on the plain client with no lock.""" + from fastapi import HTTPException + + from litellm.proxy.management_endpoints.model_management_endpoints import ( + HEURISTIC_V2_SLOT_LOCK_KEY, + _heuristic_v2_slot, + ) + + fake = self._FakeDb(db_held) + live_router = self._live_router_holding_one_heuristic_v2(limit) if config_holds_one else None + with ( + patch("litellm.proxy.proxy_server._license_check.heuristic_v2_router_limit", lambda: limit), # test-quality-ok: the guard reads the proxy license singleton with no injection seam + patch("litellm.proxy.proxy_server.llm_router", live_router), # test-quality-ok: the guard reads the proxy router global with no injection seam + patch( # test-quality-ok: the cross-pod publish is the side effect under test; redis is not configured here + "litellm.proxy.management_endpoints.model_management_endpoints.publish_config_change", + new=AsyncMock(), + ) as published, + ): + if expected == "refused": + with pytest.raises(HTTPException) as exc_info: + async with _heuristic_v2_slot(fake, effective_config=effective_config, model_id=model_id): + pass + assert exc_info.value.status_code == 403 + assert "At most 1 auto-router" in str(exc_info.value.detail) + assert "'auto_router' feature lifts the limit" in str(exc_info.value.detail) + return + async with _heuristic_v2_slot(fake, effective_config=effective_config, model_id=model_id) as tables: + handle = tables + if expected == "plain": + await handle.create(data={}) + fake.litellm_proxymodeltable.create.assert_awaited_once_with(data={}) + assert fake.tx_obj.raw_calls == [] + return + assert handle is fake.tx_obj.litellm_proxymodeltable + published.assert_awaited_once_with(redis_cache=None, object_type="litellm_proxymodeltable") + (lock_sql, lock_params), (_count_sql, count_params) = fake.tx_obj.raw_calls + assert "pg_advisory_xact_lock($1)" in lock_sql and "count" not in lock_sql + assert lock_params == (HEURISTIC_V2_SLOT_LOCK_KEY,) + assert count_params == (model_id or "",) + + @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 + (and its advisory lock) has closed; a pool-sized burst of team creates would otherwise stall on the + lock holder waiting for a connection the waiters are occupying.""" + from contextlib import asynccontextmanager + + from litellm.proxy.management_endpoints.model_management_endpoints import _add_team_model_to_db + from litellm.types.router import ModelInfo + + events: list[str] = [] + created = MagicMock(model_id="row-1") + + @asynccontextmanager + async def slot(): + events.append("slot-enter") + yield MagicMock(create=AsyncMock(return_value=created)) + events.append("slot-exit") + + async def team_model_add(**_: object) -> None: + events.append("team_model_add") + + deployment = Deployment( + model_name="public-v2", + litellm_params=LiteLLM_Params(model="auto_router/complexity_router", complexity_router_config=self._V2), + model_info=ModelInfo(id="row-1", team_id="team-1"), + ) + with ( + patch( # test-quality-ok: params are encrypted with the proxy master key, which this test does not configure + "litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper", + lambda value, new_encryption_key=None: value, + ), + patch( # test-quality-ok: the team list write is the collaborator whose ordering is asserted + "litellm.proxy.management_endpoints.model_management_endpoints.team_model_add", + side_effect=team_model_add, + ), + ): + result = await _add_team_model_to_db( + model_params=deployment, + user_api_key_dict=UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN), + prisma_client=MagicMock(), + slot=slot(), + ) + + assert result is created + assert events == ["slot-enter", "slot-exit", "team_model_add"] + + @pytest.mark.asyncio + async def test_add_new_model_refuses_a_second_heuristic_v2_router_before_the_db_write(self) -> None: + from litellm.proxy._types import ProxyException + from litellm.proxy.management_endpoints.model_management_endpoints import ( + add_new_model, + ) + + admin = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) + fake = self._FakeDb(db_held=1) + + 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.heuristic_v2_router_limit", lambda: 1), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch( # test-quality-ok: prior auth check needs a live DB; only the license limit 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-v2", + litellm_params=LiteLLM_Params(model="auto_router/complexity_router", complexity_router_config=self._V2), + ), + user_api_key_dict=admin, + ) + assert exc_info.value.code == "403" + assert "At most 1 auto-router" 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_patch_model_refuses_switching_another_router_to_heuristic_v2(self) -> None: + """patch_model relays HTTPException as-is, so the license refusal reaches the client as a plain 403.""" + from fastapi import HTTPException + + from litellm.proxy.management_endpoints.model_management_endpoints import ( + patch_model, + ) + from litellm.types.router import updateLiteLLMParams + + model_id = "other-id" + admin = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) + fake = self._FakeDb(db_held=1) + + 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.heuristic_v2_router_limit", lambda: 1), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch( # test-quality-ok: the write must be refused before this DB step runs + "litellm.proxy.management_endpoints.model_management_endpoints.get_db_model", + new=AsyncMock(return_value=self._db_complexity_router(model_id)), + ), + patch( # test-quality-ok: prior auth check needs a live DB; only the license limit 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: the helper's team bookkeeping needs a live DB; the row writer it is handed is what is under test + "litellm.proxy.management_endpoints.model_management_endpoints._update_team_model_in_db", + new=AsyncMock(side_effect=_write_empty_row), + ), + ): + with pytest.raises(HTTPException) as exc_info: + await patch_model( + model_id=model_id, + patch_data=updateDeployment(litellm_params=updateLiteLLMParams(complexity_router_config=self._V2)), + user_api_key_dict=admin, + ) + assert exc_info.value.status_code == 403 + fake.tx_obj.litellm_proxymodeltable.update.assert_not_awaited() + fake.litellm_proxymodeltable.update.assert_not_awaited() + + @pytest.mark.asyncio + async def test_update_model_refuses_switching_another_router_to_heuristic_v2(self) -> None: + from litellm.proxy._types import ProxyException + from litellm.proxy.management_endpoints.model_management_endpoints import ( + update_model, + ) + from litellm.types.router import ModelInfo, updateLiteLLMParams + + model_id = "other-id" + admin = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) + + existing_row = MagicMock() + existing_row.model_dump.return_value = { + "model_name": "my-auto-router", + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": {"tiers": {"SIMPLE": "gpt-4o-mini"}}, + }, + "model_info": {"id": model_id}, + } + existing_row.litellm_params = existing_row.model_dump.return_value["litellm_params"] + fake = self._FakeDb(db_held=1, existing_row=existing_row) + + 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.heuristic_v2_router_limit", lambda: 1), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch( # test-quality-ok: prior auth check needs a live DB; only the license limit is under test + "litellm.proxy.management_endpoints.model_management_endpoints.ModelManagementAuthChecks.can_user_make_model_call", + new=AsyncMock(return_value=None), + ), + ): + with pytest.raises(ProxyException) as exc_info: + await update_model( + model_params=updateDeployment( + litellm_params=updateLiteLLMParams(complexity_router_config=self._V2), + model_info=ModelInfo(id=model_id), + ), + user_api_key_dict=admin, + ) + assert exc_info.value.code == "403" + fake.tx_obj.litellm_proxymodeltable.update.assert_not_awaited() + fake.litellm_proxymodeltable.update.assert_not_awaited() + @pytest.mark.asyncio async def test_update_model_rejects_prefix_strip(self): from litellm.proxy._types import ProxyException diff --git a/tests/test_litellm/proxy/management_endpoints/test_ptu_model_settings.py b/tests/test_litellm/proxy/management_endpoints/test_ptu_model_settings.py index a1c38d26b9d..d35b77f732c 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ptu_model_settings.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ptu_model_settings.py @@ -45,6 +45,10 @@ from litellm.types.router import ( from litellm.types.utils import Usage +async def _passthrough_row(update_data): + return update_data + + def test_model_info_accepts_valid_ptu_fields(): info = ModelInfo( id="x", @@ -385,6 +389,7 @@ class TestTeamModelUpdateValidatesBeforeWriting: patch_data=patch_data, user_api_key_dict=MagicMock(), prisma_client=MagicMock(), + write_row=_passthrough_row, ) return result, touched @@ -914,6 +919,7 @@ class TestPtuDeploymentsAreNotBilledPerToken: patch_data=patch, user_api_key_dict=UserAPIKeyAuth(user_id="a", user_role=LitellmUserRoles.PROXY_ADMIN), prisma_client=MagicMock(), + write_row=_passthrough_row, ) assert exc.value.status_code == 400 diff --git a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py index 1ab18639fff..dcfad8f6815 100644 --- a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py +++ b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py @@ -28,6 +28,7 @@ from litellm.proxy.proxy_server import ( resolve_routing_plugins, validate_deployment_complexity_router_placement, validate_deployment_max_agentic_loops, + validate_heuristic_v2_router_limit, ) from .conftest import normalize @@ -193,6 +194,120 @@ def test_validate_deployment_complexity_router_placement_leaves_valid_deployment assert model["litellm_params"] == litellm_params +def _heuristic_v2_row(model_name: str, classifier_type: str = "heuristic_v2") -> dict[str, object]: + return { + "model_name": model_name, + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": {"classifier_type": classifier_type, "tiers": {"SIMPLE": "gpt-4o-mini"}}, + }, + } + + +def test_validate_heuristic_v2_router_limit_refuses_to_start_over_the_limit() -> None: + """Same reason as the two validators above: the proxy router swallows registration errors, so + an over-limit config.yaml must fail here instead of booting with a silently missing router.""" + with pytest.raises(ValueError, match=re.escape("At most 1 auto-router")) as exc_info: + validate_heuristic_v2_router_limit( + [_heuristic_v2_row("a"), _heuristic_v2_row("b"), _heuristic_v2_row("c", "heuristic")], limit=1 + ) + assert "'auto_router' feature lifts the limit" in str(exc_info.value) + + +@pytest.mark.parametrize( + "model_list,limit", + [ + ([_heuristic_v2_row("a"), _heuristic_v2_row("b")], None), + ([_heuristic_v2_row("a"), _heuristic_v2_row("c", "heuristic")], 1), + ([{"model_name": "gpt-4o", "litellm_params": {"model": "gpt-4o"}}], 1), + ], +) +def test_validate_heuristic_v2_router_limit_leaves_configs_within_the_limit_alone( + model_list: list[dict[str, object]], limit: int | None +) -> None: + assert validate_heuristic_v2_router_limit(model_list, limit=limit) is None + + +_TWO_HEURISTIC_V2_ROUTERS_YAML = ( + "model_list:\n" + " - model_name: gpt-4o-mini\n" + " litellm_params:\n" + " model: openai/gpt-4o-mini\n" + " api_key: k\n" + " - model_name: v2-a\n" + " litellm_params:\n" + " model: auto_router/complexity_router\n" + " complexity_router_config:\n" + " classifier_type: heuristic_v2\n" + " tiers: {SIMPLE: gpt-4o-mini}\n" + " - model_name: v2-b\n" + " litellm_params:\n" + " model: auto_router/complexity_router\n" + " complexity_router_config:\n" + " classifier_type: heuristic_v2\n" + " tiers: {SIMPLE: gpt-4o-mini}\n" + "router_settings:\n" + " heuristic_v2_router_limit: 99\n" +) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("license_limit", [1, None]) +async def test_ProxyConfig_load_config_takes_the_heuristic_v2_limit_from_the_license_only( + tmp_path, monkeypatch, license_limit: int | None +) -> None: + """`router_settings.heuristic_v2_router_limit` is managed outside config.yaml: an operator + cannot grant the entitlement by editing the config, and a licensed proxy boots both routers.""" + f = tmp_path / "c.yaml" + f.write_text(_TWO_HEURISTIC_V2_ROUTERS_YAML) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False) + monkeypatch.delenv("LITELLM_CONFIG_BUCKET_NAME", raising=False) + monkeypatch.setattr( + "litellm.proxy.proxy_server._license_check.heuristic_v2_router_limit", lambda: license_limit + ) + + if license_limit is None: + router, _model_list, _general_settings = await ProxyConfig().load_config( + router=None, config_file_path=str(f) + ) + assert router.heuristic_v2_router_limit is not None + assert router.heuristic_v2_router_limit() is None + assert sorted(router.complexity_routers) == ["v2-a", "v2-b"] + return + + with pytest.raises(ValueError, match=re.escape("config.yaml model_list: At most 1 auto-router")): + await ProxyConfig().load_config(router=None, config_file_path=str(f)) + + +@pytest.mark.asyncio +async def test_ProxyConfig_load_config_router_refuses_a_db_heuristic_v2_router_beyond_the_license( + tmp_path, monkeypatch +) -> None: + """config.yaml holds the one allowed heuristic_v2 router; a second one arriving later from the DB + is refused at registration because the router was built with the license's ceiling.""" + from litellm.types.router import Deployment + + f = tmp_path / "c.yaml" + f.write_text(_TWO_HEURISTIC_V2_ROUTERS_YAML.replace(" - model_name: v2-b\n", " - model_name: v1-b\n", 1).replace( + "classifier_type: heuristic_v2\n tiers: {SIMPLE: gpt-4o-mini}\nrouter_settings", + "classifier_type: heuristic\n tiers: {SIMPLE: gpt-4o-mini}\nrouter_settings", + )) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False) + monkeypatch.delenv("LITELLM_CONFIG_BUCKET_NAME", raising=False) + monkeypatch.setattr("litellm.proxy.proxy_server._license_check.heuristic_v2_router_limit", lambda: 1) + + router, _model_list, _general_settings = await ProxyConfig().load_config(router=None, config_file_path=str(f)) + + assert router.heuristic_v2_router_limit is not None + assert router.heuristic_v2_router_limit() == 1 + assert sorted(router.complexity_routers) == ["v1-b", "v2-a"] + db_row = Deployment(**_heuristic_v2_row("v2-from-db"), model_info={"id": "db-id"}) + assert router.upsert_deployment(db_row) is None + assert sorted(router.complexity_routers) == ["v1-b", "v2-a"] + + def test_validate_deployment_max_agentic_loops_allows_a_deployment_without_the_key(): model = {"model_name": "gpt-4o", "litellm_params": {"model": "gpt-4o"}} diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index aa1b51afe10..da3791da39a 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -14,6 +14,7 @@ from pydantic import ValidationError import litellm from litellm import Router +from litellm.router_utils.auto_router_model_naming import count_heuristic_v2_routers from litellm._logging import verbose_router_logger from litellm.caching.dual_cache import DualCache from litellm.constants import RETURN_RAW_MODEL_NAME_METADATA_KEY, SESSION_ID_GENERATED_METADATA_KEY @@ -1085,6 +1086,165 @@ class TestRouterComplexityDeploymentMethods: router.init_complexity_router_deployment(deployment) assert "auto_router/complexity_router/test-router" in router.complexity_routers + @staticmethod + def _router_row(model_name: str, model_id: str, classifier_type: str) -> dict[str, object]: + return { + "model_name": model_name, + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": { + "classifier_type": classifier_type, + "tiers": {"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o"}, + }, + }, + "model_info": {"id": model_id}, + } + + _POOL: dict[str, object] = { + "model_name": "gpt-4o-mini", + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "k"}, + } + + def test_heuristic_v2_ceiling_keeps_the_first_router_and_drops_the_rest(self) -> None: + """The proxy runs with ignore_invalid_deployments, so the second heuristic_v2 router is dropped + at registration while a heuristic (v1) sibling and the first v2 router stay routable.""" + router = Router( + model_list=[ + self._POOL, + self._router_row("v2-a", "id-a", "heuristic_v2"), + self._router_row("v2-b", "id-b", "heuristic_v2"), + self._router_row("v1-c", "id-c", "heuristic"), + ], + heuristic_v2_router_limit=lambda: 1, + ignore_invalid_deployments=True, + ) + + assert sorted(router.complexity_routers) == ["v1-c", "v2-a"] + assert router.get_deployment(model_id="id-b") is None + + def test_heuristic_v2_ceiling_raises_without_ignore_invalid_deployments(self) -> None: + with pytest.raises(ValueError, match="At most 1 auto-router"): + Router( + model_list=[ + self._POOL, + self._router_row("v2-a", "id-a", "heuristic_v2"), + self._router_row("v2-b", "id-b", "heuristic_v2"), + ], + heuristic_v2_router_limit=lambda: 1, + ) + + def test_heuristic_v2_limit_is_resolved_on_every_registration(self) -> None: + """The Router never caches the limit: when the resolver's answer moves (the proxy re-verified + its license), the next registration and the next limit query see the new value.""" + limits = {"value": None} + router = Router( + model_list=[ + self._POOL, + self._router_row("v2-a", "id-a", "heuristic_v2"), + self._router_row("v2-b", "id-b", "heuristic_v2"), + ], + heuristic_v2_router_limit=lambda: limits["value"], + ignore_invalid_deployments=True, + ) + assert sorted(router.complexity_routers) == ["v2-a", "v2-b"] + assert router.heuristic_v2_router_limit_violation() is None + + limits["value"] = 1 + assert router.heuristic_v2_router_limit_violation() is not None + assert router.upsert_deployment(Deployment(**self._router_row("v2-c", "id-c", "heuristic_v2"))) is None + assert sorted(router.complexity_routers) == ["v2-a", "v2-b"] + + def test_heuristic_v2_ceiling_tightening_refuses_the_edit_and_keeps_the_live_router(self) -> None: + """Two heuristic_v2 routers registered under an unlimited ceiling, then the ceiling drops to one: + an edit to either must be refused before its live row is popped, or the failed re-add and + the failed restore would drop a serving router while the write reports success.""" + limits = {"value": None} + router = Router( + model_list=[ + self._POOL, + self._router_row("v2-a", "id-a", "heuristic_v2"), + self._router_row("v2-b", "id-b", "heuristic_v2"), + ], + heuristic_v2_router_limit=lambda: limits["value"], + ignore_invalid_deployments=True, + ) + limits["value"] = 1 + + assert router.upsert_deployment(Deployment(**self._router_row("v2-a-renamed", "id-a", "heuristic_v2"))) is None + assert sorted(router.complexity_routers) == ["v2-a", "v2-b"] + assert router.get_deployment(model_id="id-a") is not None + + assert router.upsert_deployment(Deployment(**self._router_row("v1-a", "id-a", "heuristic"))) is not None + assert sorted(router.complexity_routers) == ["v1-a", "v2-b"] + + def test_config_deployments_excludes_db_rows(self) -> None: + """The proxy counts config.yaml routers from here and DB rows from the database, so a DB-loaded + row (``model_info.db_model``) must not show up twice.""" + router = Router(model_list=[self._POOL, self._router_row("v2-a", "id-a", "heuristic_v2")]) + db_row = self._router_row("v2-db", "id-db", "heuristic_v2") + db_row["model_info"] = {"id": "id-db", "db_model": True} + assert router.upsert_deployment(Deployment(**db_row)) is not None + + assert sorted(str(row["model_name"]) for row in router.config_deployments()) == ["gpt-4o-mini", "v2-a"] + assert count_heuristic_v2_routers(router.config_deployments()) == 1 + + def test_failed_edit_of_a_live_v2_router_rolls_back_without_the_ceiling(self) -> None: + """A rollback after a failed upsert re-admits state that was already serving, so it must not be + judged by a ceiling that tightened since: converting one of two live heuristic_v2 routers to a + config whose registration fails must leave it serving its previous v2 configuration.""" + limits = {"value": None} + router = Router( + model_list=[ + self._POOL, + self._router_row("v2-a", "id-a", "heuristic_v2"), + self._router_row("v2-b", "id-b", "heuristic_v2"), + ], + heuristic_v2_router_limit=lambda: limits["value"], + ignore_invalid_deployments=True, + ) + limits["value"] = 1 + + broken = self._router_row("v1-a", "id-a", "heuristic") + broken["litellm_params"]["complexity_router_config"]["tiers"] = {} + assert router.upsert_deployment(Deployment(**broken)) is None + + assert sorted(router.complexity_routers) == ["v2-a", "v2-b"] + live = router.get_deployment(model_id="id-a") + assert live is not None and live.litellm_params.complexity_router_config["classifier_type"] == "heuristic_v2" + assert router.heuristic_v2_router_limit_violation() is not None + + def test_heuristic_v2_routers_are_unlimited_by_default(self) -> None: + router = Router( + model_list=[ + self._POOL, + self._router_row("v2-a", "id-a", "heuristic_v2"), + self._router_row("v2-b", "id-b", "heuristic_v2"), + ] + ) + + assert sorted(router.complexity_routers) == ["v2-a", "v2-b"] + assert router.heuristic_v2_router_limit_violation() is None + + def test_heuristic_v2_router_limit_violation_frees_the_slot_of_the_router_being_edited(self) -> None: + """A DB reload upserts the existing heuristic_v2 router again; that edit must keep its own slot + while a different deployment switching to heuristic_v2 is refused.""" + router = Router( + model_list=[self._POOL, self._router_row("v2-a", "id-a", "heuristic_v2")], + heuristic_v2_router_limit=lambda: 1, + ignore_invalid_deployments=True, + ) + + assert router.heuristic_v2_router_limit_violation() is not None + + edited = self._router_row("v2-a-renamed", "id-a", "heuristic_v2") + assert router.upsert_deployment(Deployment(**edited)) is not None + assert sorted(router.complexity_routers) == ["v2-a-renamed"] + + assert router.upsert_deployment(Deployment(**self._router_row("v2-b", "id-b", "heuristic_v2"))) is None + assert sorted(router.complexity_routers) == ["v2-a-renamed"] + assert router.upsert_deployment(Deployment(**self._router_row("v1-c", "id-c", "heuristic"))) is not None + assert sorted(router.complexity_routers) == ["v1-c", "v2-a-renamed"] + def test_hybrid_initialization_waits_for_later_pool_deployments(self): router = Router( model_list=[ diff --git a/tests/test_litellm/router_utils/test_auto_router_model_naming.py b/tests/test_litellm/router_utils/test_auto_router_model_naming.py index 0007f09896a..238d0546518 100644 --- a/tests/test_litellm/router_utils/test_auto_router_model_naming.py +++ b/tests/test_litellm/router_utils/test_auto_router_model_naming.py @@ -1,8 +1,13 @@ +from collections.abc import Mapping + import pytest from litellm.router_utils.auto_router_model_naming import ( carries_complexity_router_settings, classify_strategy_router_model, + count_heuristic_v2_routers, + heuristic_v2_limit_violation, + is_heuristic_v2_router, strategy_router_dependencies, validate_complexity_router_config_placement, validate_complexity_router_config_write, @@ -369,3 +374,55 @@ def test_placement_is_scoped_to_complexity_router_deployments(model, present_fie flat param on an s3_vectors vector store, so an unscoped gate would reject a valid deployment. Either complexity field names one on its own, which is what the load itself requires.""" assert carries_complexity_router_settings(model, present_fields) is scoped + + +@pytest.mark.parametrize( + "litellm_params,expected", + [ + ({"model": "auto_router/complexity_router", "complexity_router_config": {"classifier_type": "heuristic_v2"}}, True), + ({"model": "auto_router/complexity_router-eu", "complexity_router_config": {"classifier_type": "heuristic_v2"}}, True), + ({"model": "auto_router/complexity_router", "complexity_router_config": {"classifier_type": "heuristic"}}, False), + ({"model": "auto_router/complexity_router", "complexity_router_config": {"tiers": {"SIMPLE": "a"}}}, False), + ({"model": "auto_router/complexity_router"}, False), + ({"model": "auto_router/quality_router", "complexity_router_config": {"classifier_type": "heuristic_v2"}}, False), + ({"model": "openai/gpt-4o", "complexity_router_config": {"classifier_type": "heuristic_v2"}}, False), + ({"model": "auto_router/complexity_router", "complexity_router_config": "heuristic_v2"}, False), + ({}, False), + ], +) +def test_is_heuristic_v2_router(litellm_params: Mapping[str, object], expected: bool) -> None: + """Only a complexity router whose config selects heuristic_v2 counts toward the license limit.""" + assert is_heuristic_v2_router(litellm_params) is expected + + +def test_count_heuristic_v2_routers_reads_model_list_rows_and_ignores_malformed_ones() -> None: + v2 = {"model": "auto_router/complexity_router", "complexity_router_config": {"classifier_type": "heuristic_v2"}} + rows: list[Mapping[str, object]] = [ + {"model_name": "a", "litellm_params": v2}, + {"model_name": "b", "litellm_params": {"model": "openai/gpt-4o"}}, + {"model_name": "c", "litellm_params": v2}, + {"model_name": "d"}, + {"model_name": "e", "litellm_params": "not a mapping"}, + ] + assert count_heuristic_v2_routers(rows) == 2 + assert count_heuristic_v2_routers(()) == 0 + + +@pytest.mark.parametrize( + "held,limit,violates", + [ + (1, 1, False), + (2, 1, True), + (0, 1, False), + (5, None, False), + (3, 3, False), + (4, 3, True), + ], +) +def test_heuristic_v2_limit_violation(held: int, limit: int | None, violates: bool) -> None: + violation = heuristic_v2_limit_violation(held=held, limit=limit) + assert (violation is not None) is violates + if violation is not None: + assert f"At most {limit} auto-router" in violation + assert f"would make {held}" in violation + assert "license" not in violation