scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-05 08:07:05 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 137

feat(router): limit heuristic_v2 auto-routers to one without the auto_router license feature (#39468)

Browse source 
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.
This commit is contained in:
tin-berri 2026-09-03 13:39:58 -07:00 • committed by GitHub
parent 828d561fcc
commit df73c623b2
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
13 changed files with 1067 additions and 95 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
litellm/constants.py
View file

@ -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))

23
litellm/proxy/auth/litellm_license.py
View file

@ -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,

198
litellm/proxy/management_endpoints/model_management_endpoints.py
View file

@ -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,

20
litellm/proxy/proxy_server.py
View file

@ -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:

46
litellm/router.py
View file

@ -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, ...]:

34
litellm/router_utils/auto_router_model_naming.py
View file

@ -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.

12
litellm/types/router.py
View file

@ -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

69
tests/test_litellm/proxy/auth/test_litellm_license.py
View file

@ -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

421
tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py
View file

@ -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

6
tests/test_litellm/proxy/management_endpoints/test_ptu_model_settings.py
View file

@ -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

115
tests/test_litellm/proxy/proxy_server/test_proxy_config.py
View file

@ -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"}}

160
tests/test_litellm/router_strategy/test_complexity_router.py
View file

@ -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=[

57
tests/test_litellm/router_utils/test_auto_router_model_naming.py
View file

@ -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