scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-09 22:31:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 86

feat(router): meter auto-router tier and prompt customization against the auto_router license feature (#39674)

Browse source 
Generalizes the heuristic_v2 ceiling from #39468 into a capability table whose
records own their in-process predicate, SQL spelling and refusal wording. The
existing heuristic_v2 capability keeps its own one-router ceiling. A single
customization capability combines operator-defined tier definitions with every
operator-written part of the classifier prompt. The prompt half only applies to
classifier types that call an LLM. The shipped default prompt, classification
rubric presets, tier-label renames and tier model choices remain ungated.

Scope every enforcement point to actual complexity routers. A model-less PATCH
or legacy update now decrypts the stored model before accepting strategy-router
settings, so a regular model cannot acquire a router config or spend a license
slot. Under the existing advisory lock, the cross-pod candidate query returns
only model scalars and the count decrypts and classifies them in process; old
non-router rows carrying a capability-shaped config no longer block a real
complexity router. The signed auto_router license feature makes both ceilings
unlimited.
This commit is contained in:
tin-berri 2026-09-05 09:51:23 -07:00 • committed by GitHub
parent 74ebe2974c
commit d0d09e5343
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
12 changed files with 987 additions and 205 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
litellm/constants.py
View file

@ -40,7 +40,7 @@ ROUTER_SETTINGS_MANAGED_OUTSIDE_CONFIG: Final[frozenset[str]] = frozenset(
"router_general_settings",
"ignore_invalid_deployments",
"fallback_access_check",
"heuristic_v2_router_limit",
"auto_router_capability_limit",
}
)
DEFAULT_BATCH_SIZE: Final = int(os.getenv("DEFAULT_BATCH_SIZE", 512))

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

@ -17,7 +17,7 @@ if TYPE_CHECKING:
AUTO_ROUTER_LICENSE_FEATURE: Final = "auto_router"
HEURISTIC_V2_LICENSE_REMEDY: Final = "A LiteLLM license with the 'auto_router' feature lifts the limit."
AUTO_ROUTER_LICENSE_REMEDY: Final = "A LiteLLM license with the 'auto_router' feature lifts the limit."
class LicenseCheck:
@ -153,11 +153,12 @@ class LicenseCheck:
return False
return team_count > _max_teams_in_license
def heuristic_v2_router_limit(self) -> int | None:
def auto_router_capability_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.
How many auto-routers may claim each licensed capability (heuristic_v2, operator-defined
tier_definitions): unlimited (None) only when the signed license lists the auto_router
feature, otherwise one per capability. 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

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

@ -50,7 +50,7 @@ from litellm.proxy._types import (
TeamModelDeleteRequest,
UserAPIKeyAuth,
)
from litellm.proxy.auth.litellm_license import HEURISTIC_V2_LICENSE_REMEDY
from litellm.proxy.auth.litellm_license import AUTO_ROUTER_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,
@ -98,11 +98,13 @@ from litellm.router_strategy.complexity_router import (
normalize_classification_prompt,
)
from litellm.router_utils.auto_router_model_naming import (
GATED_AUTO_ROUTER_CAPABILITIES,
STRATEGY_ROUTER_PARAM_FIELDS,
capability_limit_violation,
carries_complexity_router_settings,
count_heuristic_v2_routers,
heuristic_v2_limit_violation,
uses_heuristic_v2_classifier,
count_capability_routers,
gated_capability_of,
is_complexity_router_model,
validate_complexity_router_config_placement,
validate_complexity_router_config_write,
validate_strategy_router_model_write,
@ -237,11 +239,13 @@ def _strategy_router_write_violation(
An auto-router deployment's ``litellm_params.model`` (``auto_router/...``) is
the discriminator the router loads it by; a write that mangles it makes the
router drop the deployment silently under ``ignore_invalid_deployments``.
Only writes that supply ``litellm_params.model`` are judged on the naming
contract, against the merged (stored + incoming) params, so partial patches
and restores of an already-corrupted row stay legal. A config is judged only
when the write carries one, for the same reason: a rename must not be held
hostage by a stored config it does not touch. Returns the violation, or None.
A patch adding auto-router settings is judged against the effective model,
decrypting the stored model when the patch omits it, so a regular deployment
cannot claim a strategy-router configuration. Unrelated partial patches and
restores that do not touch strategy-router settings stay legal. A config is
judged only when the write carries one, for the same reason: a rename must
not be held hostage by a stored config it does not touch. Returns the
violation, or None.
"""
if incoming_params is None:
return None
@ -256,14 +260,18 @@ def _strategy_router_write_violation(
for source in (incoming_params, existing_params)
if source is not None and getattr(source, field, None) is not None
)
# Scope reads the incoming model because the stored one is encrypted at rest.
if carries_complexity_router_settings(incoming_params.model, present_fields):
effective_params: Final = _effective_complexity_router_params(incoming_params, existing_params)
effective_model: Final = effective_params.get("model")
if carries_complexity_router_settings(
effective_model if isinstance(effective_model, str) else None, present_fields
):
placement_violation: Final = validate_complexity_router_config_placement(incoming_params.model_extra)
if placement_violation is not None:
return placement_violation
if incoming_params.model is None:
return None
return validate_strategy_router_model_write(model=incoming_params.model, present_fields=present_fields)
return validate_strategy_router_model_write(
model=effective_model if isinstance(effective_model, str) else "",
present_fields=present_fields,
)
def _raise_on_strategy_router_write_violation(
@ -281,14 +289,23 @@ 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"
AUTO_ROUTER_CAPABILITY_SLOT_LOCK_KEY: Final = 5_872_301
_CAPABILITY_LOCK_SQL: Final = "SELECT 1 AS locked FROM pg_advisory_xact_lock($1)"
_STORED_LITELLM_PARAMS_SQL: Final = (
"(CASE jsonb_typeof(litellm_params) WHEN 'string' THEN (litellm_params #>> '{}')::jsonb ELSE litellm_params END)"
)
_STORED_COMPLEXITY_CONFIG_SQL: Final = f"{_STORED_LITELLM_PARAMS_SQL} -> 'complexity_router_config'"
_CAPABILITY_DB_ROWS_SQL: Final[Mapping[str, str]] = MappingProxyType(
{
capability.key: f"""
SELECT {_STORED_LITELLM_PARAMS_SQL} ->> 'model' AS model
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'
AND ({capability.sql_config_predicate.format(config=_STORED_COMPLEXITY_CONFIG_SQL)})
"""
for capability in GATED_AUTO_ROUTER_CAPABILITIES
}
)
def _effective_complexity_router_config(
@ -301,13 +318,44 @@ def _effective_complexity_router_config(
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.
def _effective_model(
incoming_params: GenericLiteLLMParams | None, existing_params: GenericLiteLLMParams | None
) -> str | None:
"""The model a write leaves on the row, decrypting an existing value only when the patch omits it."""
incoming: Final = None if incoming_params is None else incoming_params.model
if incoming is not None:
return incoming
existing: Final = None if existing_params is None else existing_params.model
if existing is None:
return None
decrypted: Final = decrypt_value_helper(
value=existing,
key="model",
exception_type="debug",
return_original_value=True,
)
return decrypted if isinstance(decrypted, str) else None
A write that leaves the row on classifier_type heuristic_v2 under a limited license runs
def _effective_complexity_router_params(
incoming_params: GenericLiteLLMParams | None, existing_params: GenericLiteLLMParams | None
) -> Mapping[str, object]:
"""The model and complexity config a write leaves, for placement and capability decisions."""
return MappingProxyType(
{
"model": _effective_model(incoming_params, existing_params),
"complexity_router_config": _effective_complexity_router_config(incoming_params, existing_params),
}
)
@asynccontextmanager
async def _auto_router_capability_slot(
prisma_client: PrismaClient, *, effective_params: Mapping[str, object], model_id: str | None
) -> AsyncGenerator[_ProxyModelTable, None]:
"""Hand out the model table to write through while the row's claim on a licensed capability is settled.
A write that leaves the row claiming a licensed capability 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
@ -321,21 +369,37 @@ async def _heuristic_v2_slot(
"""
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):
limit: Final = _license_check.auto_router_capability_limit()
capability: Final = gated_capability_of(effective_params)
if limit is None or capability is None:
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
await tx_ctx.query_raw(_CAPABILITY_LOCK_SQL, AUTO_ROUTER_CAPABILITY_SLOT_LOCK_KEY)
rows: Sequence[Mapping[str, object]] = await tx_ctx.query_raw(
_CAPABILITY_DB_ROWS_SQL[capability.key], model_id or ""
)
db_held: Final = sum(
1
for row in rows
for stored_model in (row.get("model"),)
if isinstance(stored_model, str)
and is_complexity_router_model(
decrypt_value_helper(
value=stored_model,
key="model",
exception_type="debug",
return_original_value=True,
)
)
)
config_rows: Final = () if llm_router is None else tuple(llm_router.config_deployments())
held: Final = (db_held if isinstance(db_held, int) else 0) + count_heuristic_v2_routers(config_rows)
violation: Final = heuristic_v2_limit_violation(held=held + 1, limit=limit)
held: Final = db_held + count_capability_routers(config_rows, capability=capability)
violation: Final = capability_limit_violation(capability=capability, held=held + 1, limit=limit)
if violation is not None:
raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN, detail=f"{violation} {HEURISTIC_V2_LICENSE_REMEDY}"
status_code=status.HTTP_403_FORBIDDEN, detail=f"{violation} {AUTO_ROUTER_LICENSE_REMEDY}"
)
yield tables.litellm_proxymodeltable
await publish_config_change(redis_cache=coordination_redis_cache(), object_type="litellm_proxymodeltable")
@ -791,6 +855,9 @@ async def patch_model(
existing_params=db_model.litellm_params,
)
effective_params: Final = _effective_complexity_router_params(
patch_data.litellm_params, db_model.litellm_params
)
requested_model_name: Final = patch_data.model_name
stored_model_name: str | None = None
@ -799,11 +866,9 @@ async def patch_model(
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(
async with _auto_router_capability_slot(
prisma_client,
effective_config=_effective_complexity_router_config(
patch_data.litellm_params, db_model.litellm_params
),
effective_params=effective_params,
model_id=model_id,
) as table:
return await table.update(where={"model_id": model_id}, data=update_data)
@ -1959,9 +2024,12 @@ async def add_new_model(
model_params=priced_model_params,
user_api_key_dict=user_api_key_dict,
prisma_client=prisma_client,
slot=_heuristic_v2_slot(
slot=_auto_router_capability_slot(
prisma_client,
effective_config=priced_model_params.litellm_params.complexity_router_config,
effective_params=_effective_complexity_router_params(
priced_model_params.litellm_params,
None,
),
model_id=priced_model_params.model_info.id,
),
)
@ -2110,6 +2178,9 @@ async def update_model(
incoming_params=model_params.litellm_params,
existing_params=deployment.litellm_params,
)
effective_params: Final = _effective_complexity_router_params(
model_params.litellm_params, deployment.litellm_params
)
# update DB
if store_model_in_db is True:
@ -2147,11 +2218,9 @@ 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}),
}
async with _heuristic_v2_slot(
async with _auto_router_capability_slot(
prisma_client,
effective_config=_effective_complexity_router_config(
model_params.litellm_params, deployment.litellm_params
),
effective_params=effective_params,
model_id=_model_id,
) as table:
model_response: Final = await table.update(

34
litellm/proxy/proxy_server.py
View file

@ -118,10 +118,11 @@ from litellm.router_utils.add_retry_fallback_headers import (
get_hidden_params_dict,
)
from litellm.router_utils.auto_router_model_naming import (
GATED_AUTO_ROUTER_CAPABILITIES,
STRATEGY_ROUTER_PARAM_FIELDS,
capability_limit_violation,
carries_complexity_router_settings,
count_heuristic_v2_routers,
heuristic_v2_limit_violation,
count_capability_routers,
validate_complexity_router_config_placement,
)
from litellm.types.utils import (
@ -303,7 +304,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 HEURISTIC_V2_LICENSE_REMEDY, LicenseCheck
from litellm.proxy.auth.litellm_license import AUTO_ROUTER_LICENSE_REMEDY, LicenseCheck
from litellm.proxy.auth.model_checks import (
expand_wildcard_deployments_for_model_info,
get_all_fallbacks,
@ -4340,17 +4341,28 @@ 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:
def validate_auto_router_capability_limits(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.
Refuse to start when config.yaml defines more auto-routers claiming a licensed capability than allowed.
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}")
violations: Final = tuple(
message
for capability in GATED_AUTO_ROUTER_CAPABILITIES
if (
message := capability_limit_violation(
capability=capability,
held=count_capability_routers(model_list, capability=capability),
limit=limit,
)
)
is not None
)
if violations:
raise ValueError(f"config.yaml model_list: {' '.join(violations)} {AUTO_ROUTER_LICENSE_REMEDY}")
def pin_complexity_router_model_id(model: dict) -> None: # mutable-ok: out-param, model_info is stamped in place
@ -5758,7 +5770,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())
validate_auto_router_capability_limits(model_list, limit=_license_check.auto_router_capability_limit())
print( # noqa: T201
"\033[32mLiteLLM: Proxy initialized with Config, Set models:\033[0m"
)
@ -5848,7 +5860,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,
auto_router_capability_limit=_license_check.auto_router_capability_limit,
)
if redis_usage_cache is not None and router.cache.redis_cache is None:
@ -6309,7 +6321,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,
auto_router_capability_limit=_license_check.auto_router_capability_limit,
)
verbose_proxy_logger.debug("updated llm_router: %s", llm_router)
else:

45
litellm/router.py
View file

@ -116,10 +116,11 @@ from litellm.router_utils.add_retry_fallback_headers import (
)
from litellm.router_utils.auto_router_model_naming import (
AUTO_ROUTER_MODEL_PREFIX,
GatedAutoRouterCapability,
capability_limit_violation,
claimed_capability,
classify_strategy_router_model,
count_heuristic_v2_routers,
heuristic_v2_limit_violation,
uses_heuristic_v2_classifier,
count_capability_routers,
)
from litellm.router_utils.batch_utils import (
_get_router_metadata_variable_name,
@ -208,6 +209,7 @@ from litellm.types.router import (
AlertingConfig,
AllowedFailsPolicy,
AssistantsTypedDict,
AutoRouterCapabilityLimit,
ConsumedRequestTagsStamp,
CredentialLiteLLMParams,
CustomRoutingStrategyBase,
@ -215,7 +217,6 @@ from litellm.types.router import (
DeploymentTypedDict,
FallbackAccessCheck,
GuardrailTypedDict,
HeuristicV2RouterLimit,
LiteLLM_Params,
MockRouterTestingParams,
ModelGroupInfo,
@ -692,7 +693,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,
auto_router_capability_limit: AutoRouterCapabilityLimit | None = None,
) -> None:
"""
Initialize the Router class with the given parameters for caching, reliability, and routing strategy.
@ -769,7 +770,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.auto_router_capability_limit = auto_router_capability_limit
self.fallback_access_check: Final = fallback_access_check
self.debug_level = debug_level
self.enable_pre_call_checks = enable_pre_call_checks
@ -8811,20 +8812,21 @@ class Router:
if not (isinstance(model_info, Mapping) and model_info.get("db_model")):
yield deployment
def heuristic_v2_router_limit_violation(self) -> str | None:
def auto_router_capability_violation(self, capability: GatedAutoRouterCapability) -> str | None:
"""
Why one more heuristic_v2 router cannot join this router, or None when it can.
Why one more router claiming ``capability`` 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.
edited first, so an edit of an existing gated router keeps its own slot. The limit is
resolved on every call through ``auto_router_capability_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)
limit: Final = self.auto_router_capability_limit() if self.auto_router_capability_limit is not None else None
others: Final = count_capability_routers(
(deployment for deployment in self.model_list if isinstance(deployment, Mapping)),
capability=capability,
)
return heuristic_v2_limit_violation(held=others + 1, limit=limit)
return capability_limit_violation(capability=capability, held=others + 1, limit=limit)
def init_complexity_router_deployment(self, deployment: Deployment):
"""
@ -8843,8 +8845,9 @@ 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()
capability: Final = claimed_capability(complexity_router_config)
if capability is not None:
limit_violation: Final = self.auto_router_capability_violation(capability)
if limit_violation is not None:
raise ValueError(limit_violation)
@ -9674,13 +9677,13 @@ class Router:
"""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
capability 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
limit_resolver: Final = self.auto_router_capability_limit
self.auto_router_capability_limit = None
try:
self.add_deployment(deployment=previous_deployment)
verbose_router_logger.info(
@ -9696,7 +9699,7 @@ class Router:
restore_error,
)
finally:
self.heuristic_v2_router_limit = limit_resolver
self.auto_router_capability_limit = limit_resolver
@staticmethod
def _backend_cost_map_keys(model: str, custom_llm_provider: str | None) -> tuple[str, ...]:

134
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 Iterable, Mapping, Sequence
from collections.abc import Callable, Iterable, Mapping, Sequence
from dataclasses import dataclass
from types import MappingProxyType
from typing import Final, Literal, TypeAlias
@ -81,6 +81,11 @@ def classify_strategy_router_model(model: str) -> StrategyRouterKind | None:
return "semantic"
def is_complexity_router_model(model: str | None) -> bool:
"""Whether ``model`` selects the complexity-router implementation."""
return classify_strategy_router_model(model or "") == "complexity"
def _named(value: object, role: StrategyRouterDependencyRole) -> tuple[StrategyRouterDependency, ...]:
"""One dependency from a scalar field, or none when it is absent or not a name."""
return (StrategyRouterDependency(value, role),) if isinstance(value, str) and value else ()
@ -168,20 +173,121 @@ def uses_heuristic_v2_classifier(complexity_router_config: object) -> bool:
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 defines_custom_tiers(complexity_router_config: object) -> bool:
"""Whether this complexity config replaces the built-in tier ladder with operator-defined tier_definitions.
Mirrors the SQL spelling on the capability record: only an actual array claims the capability,
so an explicit JSON null or a malformed value does not.
"""
return isinstance(_mapping(complexity_router_config).get("tier_definitions"), (list, tuple))
OPERATOR_CLASSIFIER_PROMPT_FIELDS: Final = ("classification_prompt", "classification_examples")
def defines_custom_classifier_prompt(complexity_router_config: object) -> bool:
"""Whether an operator wrote any part of this router's classifier prompt themselves.
Three spellings, all metered: a whole replacement prompt (``classifier_llm_config.system_prompt``),
replacement opening instructions (``classification_prompt``), and replacement calibration examples
(``classification_examples``). Choosing a shipped ``classification_rubric`` preset is not authoring.
Scoped to the classifier types that actually call an LLM, which is also where the config validator
accepts these fields: the heuristic scorers never read them.
"""
config: Final = _mapping(complexity_router_config)
if config.get("classifier_type") not in LLM_CLASSIFIER_TYPES:
return False
return _mapping(config.get("classifier_llm_config")).get("system_prompt") is not None or any(
config.get(field) is not None for field in OPERATOR_CLASSIFIER_PROMPT_FIELDS
)
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 uses_custom_tier_or_classifier_prompt(complexity_router_config: object) -> bool:
"""Whether this router replaces shipped tiers or its shipped classifier prompt."""
return defines_custom_tiers(complexity_router_config) or defines_custom_classifier_prompt(complexity_router_config)
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.
_LLM_CLASSIFIER_TYPES_SQL: Final = ", ".join(f"'{name}'" for name in sorted(LLM_CLASSIFIER_TYPES))
@dataclass(frozen=True, slots=True)
class GatedAutoRouterCapability:
"""A complexity-router capability the license meters, in every spelling an enforcement point needs.
``uses`` and ``sql_config_predicate`` answer the same question, in process and in a DB count over
stored ``litellm_params`` (``{config}`` is the caller's expression for the normalized
``complexity_router_config`` jsonb, substituted as many times as the predicate needs); they live
on one record so they cannot drift apart. ``subject`` and ``remedy`` build the shared refusal
message. A validated config claims at most one capability, and the validator is what makes that
true: tier_definitions rejects every heuristic classifier_type, and it also rejects the
classifier system_prompt, which in turn only applies to the classifier types heuristic_v2 is not.
"""
key: str
subject: str
remedy: str
uses: Callable[[object], bool]
sql_config_predicate: str
HEURISTIC_V2_CAPABILITY: Final = GatedAutoRouterCapability(
key="heuristic_v2",
subject="with classifier_type 'heuristic_v2'",
remedy="Use classifier_type 'heuristic' for this router or remove an existing heuristic_v2 router.",
uses=uses_heuristic_v2_classifier,
sql_config_predicate="{config} ->> 'classifier_type' = 'heuristic_v2'",
)
_OPERATOR_PROMPT_FIELDS_SQL: Final = " OR ".join(
f"{{config}} ->> '{field}' IS NOT NULL" for field in OPERATOR_CLASSIFIER_PROMPT_FIELDS
)
CUSTOMIZATION_CAPABILITY: Final = GatedAutoRouterCapability(
key="tier_or_classifier_prompt",
subject="with operator-defined tier_definitions or an operator-written classifier prompt",
remedy=(
"Use the shipped tiers and classifier prompt for this router or remove an existing router "
"with tier_definitions or its own classifier prompt."
),
uses=uses_custom_tier_or_classifier_prompt,
sql_config_predicate=(
"jsonb_typeof({config} -> 'tier_definitions') = 'array' OR "
f"({{config}} ->> 'classifier_type' IN ({_LLM_CLASSIFIER_TYPES_SQL}) AND ("
"{config} -> 'classifier_llm_config' ->> 'system_prompt' IS NOT NULL OR "
f"{_OPERATOR_PROMPT_FIELDS_SQL}))"
),
)
GATED_AUTO_ROUTER_CAPABILITIES: Final = (HEURISTIC_V2_CAPABILITY, CUSTOMIZATION_CAPABILITY)
def claimed_capability(complexity_router_config: object) -> GatedAutoRouterCapability | None:
"""The licensed capability this complexity config claims, or None."""
return next(
(capability for capability in GATED_AUTO_ROUTER_CAPABILITIES if capability.uses(complexity_router_config)),
None,
)
def gated_capability_of(litellm_params: Mapping[str, object]) -> GatedAutoRouterCapability | None:
"""The licensed capability this deployment claims, or None unless it is a complexity router."""
model: Final = litellm_params.get("model")
if not is_complexity_router_model(model if isinstance(model, str) else None):
return None
return claimed_capability(litellm_params.get("complexity_router_config"))
def count_capability_routers(
deployments: Iterable[Mapping[str, object]], *, capability: GatedAutoRouterCapability
) -> int:
"""How many of ``deployments`` (router model_list entries or config.yaml rows) claim ``capability``."""
return sum(
1 for deployment in deployments if gated_capability_of(_mapping(deployment.get("litellm_params"))) is capability
)
def capability_limit_violation(*, capability: GatedAutoRouterCapability, held: int, limit: int | None) -> str | None:
"""Why holding ``held`` routers claiming ``capability`` 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
@ -190,8 +296,8 @@ def heuristic_v2_limit_violation(*, held: int, limit: int | None) -> str | None:
if limit is None or held <= limit:
return None
return (
f"At most {limit} auto-router(s) with 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."
f"At most {limit} auto-router(s) {capability.subject} can be registered but this would make "
f"{held}. {capability.remedy}"
)
@ -237,9 +343,7 @@ def carries_complexity_router_settings(model: str | None, present_fields: frozen
``validate_strategy_router_model_write`` is judged on, so a router named only by its
default model is in scope, and a field added to the table above is covered here for free.
"""
return classify_strategy_router_model(model or "") == "complexity" or bool(
present_fields & _COMPLEXITY_ROUTER_FIELDS
)
return is_complexity_router_model(model) or bool(present_fields & _COMPLEXITY_ROUTER_FIELDS)
def validate_complexity_router_config_placement(litellm_params: Mapping[str, object] | None) -> str | None:

4
litellm/types/router.py
View file

@ -887,9 +887,9 @@ class FallbackAccessCheck(Protocol):
async def __call__(self, *, model: str, request_kwargs: Mapping[str, object], llm_router: "Router") -> bool: ...
class HeuristicV2RouterLimit(Protocol):
class AutoRouterCapabilityLimit(Protocol):
"""
Resolves how many heuristic_v2 complexity routers the Router may hold right now; None means unlimited.
Resolves how many complexity routers may claim each licensed capability 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

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

@ -34,27 +34,27 @@ def test_is_over_limit():
assert license_check.is_over_limit(99) is False
def test_heuristic_v2_router_limit() -> None:
def test_auto_router_capability_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
assert license_check.auto_router_capability_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
assert license_check.auto_router_capability_limit() is None
license_check.airgapped_license_data = {"expiration_date": "2999-01-01", "allowed_features": ["sso"]}
assert license_check.heuristic_v2_router_limit() == 1
assert license_check.auto_router_capability_limit() == 1
license_check.airgapped_license_data = {"expiration_date": "2999-01-01"}
assert license_check.heuristic_v2_router_limit() == 1
assert license_check.auto_router_capability_limit() == 1
license_check.airgapped_license_data = None
assert license_check.heuristic_v2_router_limit() == 1
assert license_check.auto_router_capability_limit() == 1
def _signed_license(expiration_date: str) -> tuple[RSAPublicKey, str]:
@ -81,12 +81,12 @@ def test_expired_or_unreadable_license_grants_no_features() -> None:
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
assert license_check.auto_router_capability_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.auto_router_capability_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
@ -98,4 +98,4 @@ def test_valid_signed_license_with_auto_router_lifts_the_limit() -> None:
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
assert license_check.auto_router_capability_limit() is None

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

@ -2,6 +2,7 @@ import inspect
import asyncio
import contextlib
import json
from collections.abc import Mapping
from typing import Dict, Optional
from unittest.mock import AsyncMock, MagicMock, patch
@ -4048,6 +4049,72 @@ class TestStrategyRouterWriteValidation:
)
assert _strategy_router_write_violation(incoming_params=None, existing_params=None) is None
@pytest.mark.parametrize(
"config",
[
{"classifier_type": "heuristic_v2", "tiers": {"SIMPLE": "gpt-4o-mini"}},
{
"classifier_type": "llm",
"classifier_llm_config": {"model": "gpt-4o-mini"},
"tier_definitions": [
{"name": "routine", "description": "routine drafting"},
{"name": "hard", "description": "hard reasoning"},
],
"tiers": {"routine": "gpt-4o-mini", "hard": "gpt-4o"},
"fallback_tier": "routine",
},
],
)
def test_model_less_patch_cannot_attach_router_config_to_a_regular_model(self, config: dict[str, object]) -> None:
"""The license gate applies only to complexity routers, so a partial PATCH cannot poison a regular
model with a capability-shaped config and make it occupy a slot."""
from litellm.proxy.management_endpoints.model_management_endpoints import (
_strategy_router_write_violation,
)
from litellm.types.router import updateLiteLLMParams
violation = _strategy_router_write_violation(
incoming_params=updateLiteLLMParams(complexity_router_config=config),
existing_params=LiteLLM_Params(model="openai/gpt-4o-mini"),
)
assert violation is not None
assert "does not start with 'auto_router/'" in violation
assert "complexity_router_config" in violation
def test_effective_params_decrypts_a_stored_complexity_router_model(self, monkeypatch) -> None:
"""A database row encrypts model, so the model-aware gate must not accidentally rely on plaintext mocks."""
from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper
from litellm.proxy.management_endpoints.model_management_endpoints import (
_effective_complexity_router_params,
)
from litellm.types.router import updateLiteLLMParams
monkeypatch.setenv("LITELLM_SALT_KEY", "test-salt")
encrypted_model = encrypt_value_helper("auto_router/complexity_router")
effective_params = _effective_complexity_router_params(
updateLiteLLMParams(complexity_router_config={"tiers": {"SIMPLE": "gpt-4o-mini"}}),
LiteLLM_Params(model=encrypted_model),
)
assert effective_params["model"] == "auto_router/complexity_router"
def test_model_less_patch_keeps_a_complexity_router_in_scope(self) -> None:
from litellm.proxy.management_endpoints.model_management_endpoints import (
_strategy_router_write_violation,
)
from litellm.types.router import updateLiteLLMParams
assert (
_strategy_router_write_violation(
incoming_params=updateLiteLLMParams(
complexity_router_config={"classifier_type": "heuristic_v2", "tiers": {"SIMPLE": "gpt-4o-mini"}}
),
existing_params=self._stored_complexity_params(),
)
is None
)
def test_restore_of_corrupted_row_is_allowed(self):
from litellm.proxy.management_endpoints.model_management_endpoints import (
_strategy_router_write_violation,
@ -4354,33 +4421,33 @@ class TestStrategyRouterWriteValidation:
)
@staticmethod
def _live_router_holding_one_heuristic_v2(limit: int | None) -> Router:
def _live_router_holding_one_capability(limit: int | None, config: Mapping[str, object]) -> Router:
return Router(
model_list=[
{"model_name": "gpt-4o-mini", "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "k"}},
{
"model_name": "held-v2",
"model_name": "held",
"litellm_params": {
"model": "auto_router/complexity_router",
"complexity_router_config": {"classifier_type": "heuristic_v2", "tiers": {"SIMPLE": "gpt-4o-mini"}},
"complexity_router_config": config,
},
"model_info": {"id": "held-id"},
},
],
heuristic_v2_router_limit=lambda: limit,
auto_router_capability_limit=lambda: limit,
)
class _FakeTx:
"""Stands in for a prisma transaction: records the raw statements and exposes the model table."""
"""Stands in for a prisma transaction: records raw statements and returns encrypted-model candidates."""
def __init__(self, db_held: int) -> None:
self.db_held = db_held
def __init__(self, db_models: list[str]) -> None:
self.db_models = db_models
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 []
return [{"model": model} for model in self.db_models] if "AS model" in sql else []
async def __aenter__(self) -> "TestStrategyRouterWriteValidation._FakeTx":
return self
@ -4391,9 +4458,9 @@ class TestStrategyRouterWriteValidation:
class _FakeDb:
"""Stands in for prisma_client: the plain client and the transaction it opens are told apart by identity."""
def __init__(self, db_held: int, existing_row: object = None) -> None:
def __init__(self, db_models: list[str], existing_row: object = None) -> None:
self.db = self
self.tx_obj = TestStrategyRouterWriteValidation._FakeTx(db_held)
self.tx_obj = TestStrategyRouterWriteValidation._FakeTx(db_models)
self.litellm_proxymodeltable = MagicMock(
create=AsyncMock(), update=AsyncMock(), find_unique=AsyncMock(return_value=existing_row)
)
@ -4403,6 +4470,43 @@ class TestStrategyRouterWriteValidation:
_V2 = {"classifier_type": "heuristic_v2", "tiers": {"SIMPLE": "gpt-4o-mini"}}
_V1 = {"classifier_type": "heuristic", "tiers": {"SIMPLE": "gpt-4o-mini"}}
_CUSTOM_TIERS = {
"classifier_type": "llm",
"classifier_llm_config": {"model": "gpt-4o-mini"},
"tier_definitions": [
{"name": "routine", "description": "routine drafting"},
{"name": "hard", "description": "hard reasoning"},
],
"tiers": {"routine": "gpt-4o-mini", "hard": "gpt-4o"},
"fallback_tier": "routine",
}
_TIER_LABELS_ONLY = {
"classifier_type": "heuristic",
"tiers": {"SIMPLE": "gpt-4o-mini"},
"tier_labels": {"SIMPLE": "Cheap"},
}
_CUSTOM_PROMPT = {
"classifier_type": "llm",
"classifier_llm_config": {"model": "gpt-4o-mini", "system_prompt": "judge it my way"},
"tiers": {"SIMPLE": "gpt-4o-mini"},
}
_OPERATOR_EXAMPLES = {
"classifier_type": "llm",
"classifier_llm_config": {"model": "gpt-4o-mini"},
"tiers": {"SIMPLE": "gpt-4o-mini"},
"classification_examples": '- "reset my password" -> SIMPLE',
}
_OPERATOR_OPENING_PROMPT = {
"classifier_type": "llm",
"classifier_llm_config": {"model": "gpt-4o-mini"},
"tiers": {"SIMPLE": "gpt-4o-mini"},
"classification_prompt": "Grade by data sensitivity",
}
_SHIPPED_RUBRIC = {
"classifier_type": "llm",
"classifier_llm_config": {"model": "gpt-4o-mini", "classification_rubric": "agentic"},
"tiers": {"SIMPLE": "gpt-4o-mini"},
}
@pytest.mark.parametrize(
"incoming,existing,expected",
@ -4431,41 +4535,55 @@ class TestStrategyRouterWriteValidation:
@pytest.mark.asyncio
@pytest.mark.parametrize(
"limit,effective_config,db_held,config_holds_one,model_id,expected",
"limit,effective_params,db_models,config_config,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"),
(1, {"model": "auto_router/complexity_router", "complexity_router_config": _V2}, ["auto_router/complexity_router"], None, None, "refused"),
(1, {"model": "auto_router/complexity_router", "complexity_router_config": _V2}, [], _V2, None, "refused"),
(1, {"model": "auto_router/complexity_router", "complexity_router_config": _V2}, [], None, None, "reserved"),
(1, {"model": "auto_router/complexity_router", "complexity_router_config": _V2}, [], None, "held-id", "reserved"),
(2, {"model": "auto_router/complexity_router", "complexity_router_config": _V2}, ["auto_router/complexity_router"], None, None, "reserved"),
(1, {"model": "auto_router/complexity_router", "complexity_router_config": _CUSTOM_TIERS}, ["openai/gpt-4o"], None, None, "reserved"),
(1, {"model": "auto_router/complexity_router", "complexity_router_config": _CUSTOM_TIERS}, [], _CUSTOM_PROMPT, None, "refused"),
(1, {"model": "openai/gpt-4o", "complexity_router_config": _CUSTOM_TIERS}, ["auto_router/complexity_router"], None, None, "plain"),
(1, {"model": "auto_router/complexity_router", "complexity_router_config": _V1}, ["auto_router/complexity_router"], _V2, None, "plain"),
(1, {"model": "auto_router/complexity_router", "complexity_router_config": None}, ["auto_router/complexity_router"], _V2, None, "plain"),
(None, {"model": "auto_router/complexity_router", "complexity_router_config": _V2}, ["auto_router/complexity_router"], _V2, None, "plain"),
(1, {"model": "auto_router/complexity_router", "complexity_router_config": _TIER_LABELS_ONLY}, ["auto_router/complexity_router"], _CUSTOM_TIERS, None, "plain"),
(1, {"model": "auto_router/complexity_router", "complexity_router_config": _CUSTOM_PROMPT}, ["auto_router/complexity_router"], None, None, "refused"),
(1, {"model": "auto_router/complexity_router", "complexity_router_config": _OPERATOR_EXAMPLES}, [], _CUSTOM_TIERS, None, "refused"),
(1, {"model": "auto_router/complexity_router", "complexity_router_config": _OPERATOR_OPENING_PROMPT}, [], _CUSTOM_PROMPT, None, "refused"),
(None, {"model": "auto_router/complexity_router", "complexity_router_config": _OPERATOR_EXAMPLES}, ["auto_router/complexity_router"], _CUSTOM_TIERS, None, "plain"),
],
)
async def test_heuristic_v2_slot_matrix(
async def test_auto_router_capability_slot_matrix(
self,
limit: int | None,
effective_config: object,
db_held: int,
config_holds_one: bool,
effective_params: Mapping[str, object],
db_models: list[str],
config_config: Mapping[str, object] | None,
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."""
"""The slot is claimed inside a locked transaction only for a write that claims a licensed capability
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.
heuristic_v2 has its own slot, while custom tier definitions and custom prompts count into one shared
customization slot. Renaming built-in tiers through tier_labels claims nothing at all."""
from fastapi import HTTPException
from litellm.proxy.management_endpoints.model_management_endpoints import (
HEURISTIC_V2_SLOT_LOCK_KEY,
_heuristic_v2_slot,
AUTO_ROUTER_CAPABILITY_SLOT_LOCK_KEY,
_auto_router_capability_slot,
)
from litellm.router_utils.auto_router_model_naming import gated_capability_of
fake = self._FakeDb(db_held)
live_router = self._live_router_holding_one_heuristic_v2(limit) if config_holds_one else None
capability = gated_capability_of(effective_params)
fake = self._FakeDb(db_models)
live_router = self._live_router_holding_one_capability(limit, config_config) if config_config is not None 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._license_check.auto_router_capability_limit", lambda: limit), # test-quality-ok: the guard reads the proxy license singleton with no injection seam
patch("litellm.proxy.proxy_server.llm_router", 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",
@ -4474,13 +4592,15 @@ class TestStrategyRouterWriteValidation:
):
if expected == "refused":
with pytest.raises(HTTPException) as exc_info:
async with _heuristic_v2_slot(fake, effective_config=effective_config, model_id=model_id):
async with _auto_router_capability_slot(fake, effective_params=effective_params, model_id=model_id):
pass
assert exc_info.value.status_code == 403
assert capability is not None
assert "At most 1 auto-router" in str(exc_info.value.detail)
assert capability.subject 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:
async with _auto_router_capability_slot(fake, effective_params=effective_params, model_id=model_id) as tables:
handle = tables
if expected == "plain":
await handle.create(data={})
@ -4489,10 +4609,13 @@ class TestStrategyRouterWriteValidation:
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
(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 lock_params == (AUTO_ROUTER_CAPABILITY_SLOT_LOCK_KEY,)
assert count_params == (model_id or "",)
assert "AS model" in count_sql
assert capability is not None
assert capability.sql_config_predicate.split("{config}")[-1].strip() in count_sql
@pytest.mark.asyncio
async def test_team_model_bookkeeping_runs_after_the_slot_is_released(self) -> None:
@ -4549,14 +4672,14 @@ class TestStrategyRouterWriteValidation:
)
admin = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN)
fake = self._FakeDb(db_held=1)
fake = self._FakeDb(["auto_router/complexity_router"])
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("litellm.proxy.proxy_server._license_check.auto_router_capability_limit", lambda: 1), # test-quality-ok: endpoint reads proxy server globals with no injection seam
patch( # 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),
@ -4579,6 +4702,93 @@ class TestStrategyRouterWriteValidation:
fake.tx_obj.litellm_proxymodeltable.create.assert_not_awaited()
fake.litellm_proxymodeltable.create.assert_not_awaited()
@pytest.mark.asyncio
async def test_model_less_patch_rejects_router_config_on_a_regular_model(self) -> None:
"""PATCH rejects the poison before its row write or the capability slot."""
from litellm.proxy._types import ProxyException
from litellm.proxy.management_endpoints.model_management_endpoints import patch_model
from litellm.types.router import updateLiteLLMParams
model_id = "regular-model"
admin = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN)
regular = Deployment(
model_name="regular-model",
litellm_params=LiteLLM_Params(model="openai/gpt-4o-mini"),
model_info={"id": model_id},
)
fake = self._FakeDb([])
with (
patch("litellm.proxy.proxy_server.prisma_client", fake), # test-quality-ok: endpoint reads proxy globals with no injection seam
patch("litellm.proxy.proxy_server.llm_router", None), # test-quality-ok: config-based lookup must be absent to drive the stored-row branch
patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: endpoint reaches its DB-write branch only with this process setting
patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: authorization branch reads the proxy-wide premium flag
patch( # test-quality-ok: inject stored regular row without a database
"litellm.proxy.management_endpoints.model_management_endpoints.get_db_model",
new=AsyncMock(return_value=regular),
),
patch( # test-quality-ok: endpoint must reject before database authorization needs a live store
"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 patch_model(
model_id=model_id,
patch_data=updateDeployment(
litellm_params=updateLiteLLMParams(complexity_router_config=self._CUSTOM_TIERS)
),
user_api_key_dict=admin,
)
assert exc_info.value.code == "400"
assert "does not start with 'auto_router/'" in str(exc_info.value.message)
assert fake.tx_obj.raw_calls == []
assert fake.tx_obj.litellm_proxymodeltable.update.await_count == 0
assert fake.litellm_proxymodeltable.update.await_count == 0
@pytest.mark.asyncio
async def test_model_less_legacy_update_rejects_router_config_on_a_regular_model(self) -> None:
"""The legacy update endpoint enforces the same boundary before its row write or slot."""
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 = "regular-model"
admin = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN)
regular = Deployment(
model_name="regular-model",
litellm_params=LiteLLM_Params(model="openai/gpt-4o-mini"),
model_info={"id": model_id},
)
existing_row = MagicMock()
existing_row.model_dump.return_value = regular.model_dump()
existing_row.litellm_params = regular.litellm_params.model_dump()
fake = self._FakeDb([], existing_row=existing_row)
with (
patch("litellm.proxy.proxy_server.prisma_client", fake), # test-quality-ok: endpoint reads proxy globals with no injection seam
patch("litellm.proxy.proxy_server.llm_router", None), # test-quality-ok: config-based lookup must be absent to drive the stored-row branch
patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: endpoint reaches its DB-write branch only with this process setting
patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: authorization branch reads the proxy-wide premium flag
patch( # test-quality-ok: endpoint must reject before database authorization needs a live store
"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._CUSTOM_TIERS),
model_info=ModelInfo(id=model_id),
),
user_api_key_dict=admin,
)
assert exc_info.value.code == "400"
assert "does not start with 'auto_router/'" in str(exc_info.value.message)
assert fake.tx_obj.raw_calls == []
assert fake.tx_obj.litellm_proxymodeltable.update.await_count == 0
assert fake.litellm_proxymodeltable.update.await_count == 0
@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."""
@ -4591,14 +4801,14 @@ class TestStrategyRouterWriteValidation:
model_id = "other-id"
admin = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN)
fake = self._FakeDb(db_held=1)
fake = self._FakeDb(["auto_router/complexity_router"])
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("litellm.proxy.proxy_server._license_check.auto_router_capability_limit", lambda: 1), # test-quality-ok: endpoint reads proxy server globals with no injection seam
patch( # 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)),
@ -4643,14 +4853,14 @@ class TestStrategyRouterWriteValidation:
"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)
fake = self._FakeDb(["auto_router/complexity_router"], 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("litellm.proxy.proxy_server._license_check.auto_router_capability_limit", lambda: 1), # test-quality-ok: endpoint reads proxy server globals with no injection seam
patch( # 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),

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

@ -28,7 +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,
validate_auto_router_capability_limits,
)
from .conftest import normalize
@ -204,13 +204,71 @@ def _heuristic_v2_row(model_name: str, classifier_type: str = "heuristic_v2") ->
}
def test_validate_heuristic_v2_router_limit_refuses_to_start_over_the_limit() -> None:
def _custom_tier_row(model_name: str) -> dict[str, object]:
return {
"model_name": model_name,
"litellm_params": {
"model": "auto_router/complexity_router",
"complexity_router_config": {
"classifier_type": "llm",
"tier_definitions": [
{"name": "routine", "description": "routine drafting"},
{"name": "hard", "description": "hard reasoning"},
],
"tiers": {"routine": "gpt-4o-mini", "hard": "gpt-4o"},
"fallback_tier": "routine",
},
},
}
def _operator_examples_row(model_name: str) -> dict[str, object]:
return {
"model_name": model_name,
"litellm_params": {
"model": "auto_router/complexity_router",
"complexity_router_config": {
"classifier_type": "llm",
"classifier_llm_config": {"model": "gpt-4o-mini"},
"tiers": {"SIMPLE": "gpt-4o-mini"},
"classification_examples": '- "reset my password" -> SIMPLE',
},
},
}
def _custom_prompt_row(model_name: str) -> dict[str, object]:
return {
"model_name": model_name,
"litellm_params": {
"model": "auto_router/complexity_router",
"complexity_router_config": {
"classifier_type": "llm",
"classifier_llm_config": {"model": "gpt-4o-mini", "system_prompt": "judge it my way"},
"tiers": {"SIMPLE": "gpt-4o-mini"},
},
},
}
@pytest.mark.parametrize(
"over_limit_rows,subject",
[
([_heuristic_v2_row("a"), _heuristic_v2_row("b"), _heuristic_v2_row("c", "heuristic")], "heuristic_v2"),
([_custom_tier_row("a"), _custom_tier_row("b"), _heuristic_v2_row("c", "heuristic")], "tier_definitions"),
([_custom_prompt_row("a"), _custom_prompt_row("b"), _heuristic_v2_row("c", "heuristic")], "operator-written classifier prompt"),
([_custom_tier_row("a"), _custom_prompt_row("b"), _heuristic_v2_row("c", "heuristic")], "operator-written classifier prompt"),
([_operator_examples_row("a"), _custom_tier_row("b"), _heuristic_v2_row("c", "heuristic")], "operator-written classifier prompt"),
],
)
def test_validate_auto_router_capability_limits_refuses_to_start_over_the_limit(
over_limit_rows: list[dict[str, object]], subject: str
) -> 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
)
validate_auto_router_capability_limits(over_limit_rows, limit=1)
assert subject in str(exc_info.value)
assert "'auto_router' feature lifts the limit" in str(exc_info.value)
@ -220,12 +278,15 @@ def test_validate_heuristic_v2_router_limit_refuses_to_start_over_the_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),
([_custom_tier_row("a"), _custom_tier_row("b")], None),
([_custom_tier_row("a"), _heuristic_v2_row("b")], 1),
],
)
def test_validate_heuristic_v2_router_limit_leaves_configs_within_the_limit_alone(
def test_validate_auto_router_capability_limits_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
"""The last case is the separate-ceiling invariant: one router of each capability fits under a limit of one."""
assert validate_auto_router_capability_limits(model_list, limit=limit) is None
_TWO_HEURISTIC_V2_ROUTERS_YAML = (
@ -247,7 +308,7 @@ _TWO_HEURISTIC_V2_ROUTERS_YAML = (
" classifier_type: heuristic_v2\n"
" tiers: {SIMPLE: gpt-4o-mini}\n"
"router_settings:\n"
" heuristic_v2_router_limit: 99\n"
" auto_router_capability_limit: 99\n"
)
@ -256,7 +317,7 @@ _TWO_HEURISTIC_V2_ROUTERS_YAML = (
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
"""`router_settings.auto_router_capability_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)
@ -264,15 +325,15 @@ async def test_ProxyConfig_load_config_takes_the_heuristic_v2_limit_from_the_lic
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
"litellm.proxy.proxy_server._license_check.auto_router_capability_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 router.auto_router_capability_limit is not None
assert router.auto_router_capability_limit() is None
assert sorted(router.complexity_routers) == ["v2-a", "v2-b"]
return
@ -296,12 +357,12 @@ async def test_ProxyConfig_load_config_router_refuses_a_db_heuristic_v2_router_b
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)
monkeypatch.setattr("litellm.proxy.proxy_server._license_check.auto_router_capability_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 router.auto_router_capability_limit is not None
assert router.auto_router_capability_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

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

@ -15,6 +15,12 @@ from pydantic import ValidationError
import litellm
from litellm import Router
from litellm.router_utils.auto_router_model_naming import (
CUSTOMIZATION_CAPABILITY,
GATED_AUTO_ROUTER_CAPABILITIES,
HEURISTIC_V2_CAPABILITY,
count_capability_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
@ -46,7 +52,6 @@ from litellm.router_strategy.complexity_router.tier_predictor import (
TierGlobalStatistic,
TrainedTierArtifact,
)
from litellm.router_utils.auto_router_model_naming import count_heuristic_v2_routers
from litellm.types.router import (
Deployment,
LiteLLM_Params,
@ -1124,7 +1129,7 @@ class TestRouterComplexityDeploymentMethods:
self._router_row("v2-b", "id-b", "heuristic_v2"),
self._router_row("v1-c", "id-c", "heuristic"),
],
heuristic_v2_router_limit=lambda: 1,
auto_router_capability_limit=lambda: 1,
ignore_invalid_deployments=True,
)
@ -1139,7 +1144,7 @@ class TestRouterComplexityDeploymentMethods:
self._router_row("v2-a", "id-a", "heuristic_v2"),
self._router_row("v2-b", "id-b", "heuristic_v2"),
],
heuristic_v2_router_limit=lambda: 1,
auto_router_capability_limit=lambda: 1,
)
def test_heuristic_v2_limit_is_resolved_on_every_registration(self) -> None:
@ -1152,14 +1157,14 @@ class TestRouterComplexityDeploymentMethods:
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"],
auto_router_capability_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
assert router.auto_router_capability_violation(HEURISTIC_V2_CAPABILITY) is None
limits["value"] = 1
assert router.heuristic_v2_router_limit_violation() is not None
assert router.auto_router_capability_violation(HEURISTIC_V2_CAPABILITY) 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"]
@ -1174,7 +1179,7 @@ class TestRouterComplexityDeploymentMethods:
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"],
auto_router_capability_limit=lambda: limits["value"],
ignore_invalid_deployments=True,
)
limits["value"] = 1
@ -1195,7 +1200,7 @@ class TestRouterComplexityDeploymentMethods:
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
assert count_capability_routers(router.config_deployments(), capability=HEURISTIC_V2_CAPABILITY) == 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
@ -1208,7 +1213,7 @@ class TestRouterComplexityDeploymentMethods:
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"],
auto_router_capability_limit=lambda: limits["value"],
ignore_invalid_deployments=True,
)
limits["value"] = 1
@ -1220,7 +1225,7 @@ class TestRouterComplexityDeploymentMethods:
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
assert router.auto_router_capability_violation(HEURISTIC_V2_CAPABILITY) is not None
def test_heuristic_v2_routers_are_unlimited_by_default(self) -> None:
router = Router(
@ -1232,18 +1237,18 @@ class TestRouterComplexityDeploymentMethods:
)
assert sorted(router.complexity_routers) == ["v2-a", "v2-b"]
assert router.heuristic_v2_router_limit_violation() is None
assert router.auto_router_capability_violation(HEURISTIC_V2_CAPABILITY) is None
def test_heuristic_v2_router_limit_violation_frees_the_slot_of_the_router_being_edited(self) -> None:
def test_auto_router_capability_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,
auto_router_capability_limit=lambda: 1,
ignore_invalid_deployments=True,
)
assert router.heuristic_v2_router_limit_violation() is not None
assert router.auto_router_capability_violation(HEURISTIC_V2_CAPABILITY) is not None
edited = self._router_row("v2-a-renamed", "id-a", "heuristic_v2")
assert router.upsert_deployment(Deployment(**edited)) is not None
@ -1254,6 +1259,205 @@ class TestRouterComplexityDeploymentMethods:
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"]
@staticmethod
def _custom_tier_row(model_name: str, model_id: str) -> dict[str, object]:
return {
"model_name": model_name,
"litellm_params": {
"model": "auto_router/complexity_router",
"complexity_router_default_model": "gpt-4o-mini",
"complexity_router_config": {
"classifier_type": "llm",
"classifier_llm_config": {"model": "gpt-4o-mini"},
"tier_definitions": [
{"name": "routine", "description": "routine drafting and lookups"},
{"name": "hard", "description": "multi-step reasoning under tradeoffs"},
],
"tiers": {"routine": "gpt-4o-mini", "hard": "gpt-4o"},
"fallback_tier": "routine",
},
},
"model_info": {"id": model_id},
}
@staticmethod
def _custom_prompt_row(model_name: str, model_id: str) -> dict[str, object]:
return {
"model_name": model_name,
"litellm_params": {
"model": "auto_router/complexity_router",
"complexity_router_default_model": "gpt-4o-mini",
"complexity_router_config": {
"classifier_type": "llm",
"classifier_llm_config": {"model": "gpt-4o-mini", "system_prompt": "judge it my way"},
"tiers": {"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o"},
},
},
} | {"model_info": {"id": model_id}}
def test_a_second_custom_prompt_router_is_refused_under_the_ceiling(self) -> None:
"""An operator-written classifier system_prompt is metered like the other licensed capabilities."""
with pytest.raises(ValueError, match="operator-written classifier prompt"):
Router(
model_list=[
self._POOL,
self._custom_prompt_row("prompt-a", "id-a"),
self._custom_prompt_row("prompt-b", "id-b"),
],
auto_router_capability_limit=lambda: 1,
)
def test_the_shipped_rubric_and_default_prompt_stay_free(self) -> None:
"""Only an operator-written prompt is gated: picking a shipped rubric preset, or writing no
prompt at all, leaves a router unmetered, so several of them register under a ceiling of one."""
def rubric(model_name: str, model_id: str, preset: str | None) -> dict[str, object]:
llm_config: dict[str, object] = {"model": "gpt-4o-mini"}
if preset is not None:
llm_config["classification_rubric"] = preset
return {
"model_name": model_name,
"litellm_params": {
"model": "auto_router/complexity_router",
"complexity_router_default_model": "gpt-4o-mini",
"complexity_router_config": {
"classifier_type": "llm",
"classifier_llm_config": llm_config,
"tiers": {"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o"},
},
},
"model_info": {"id": model_id},
}
router = Router(
model_list=[
self._POOL,
rubric("default-a", "id-a", None),
rubric("preset-b", "id-b", "agentic"),
rubric("preset-c", "id-c", "chat"),
],
auto_router_capability_limit=lambda: 1,
)
assert sorted(router.complexity_routers) == ["default-a", "preset-b", "preset-c"]
def test_a_second_custom_tier_router_is_refused_under_the_ceiling(self) -> None:
"""Operator-defined tier sets are metered like heuristic_v2: one per proxy without the license."""
with pytest.raises(ValueError, match="tier_definitions"):
Router(
model_list=[
self._POOL,
self._custom_tier_row("tiers-a", "id-a"),
self._custom_tier_row("tiers-b", "id-b"),
],
auto_router_capability_limit=lambda: 1,
)
def test_custom_tier_routers_are_unlimited_with_the_license_feature(self) -> None:
router = Router(
model_list=[
self._POOL,
self._custom_tier_row("tiers-a", "id-a"),
self._custom_tier_row("tiers-b", "id-b"),
],
auto_router_capability_limit=lambda: None,
)
assert sorted(router.complexity_routers) == ["tiers-a", "tiers-b"]
assert router.auto_router_capability_violation(CUSTOMIZATION_CAPABILITY) is None
def test_each_capability_holds_its_own_slot(self) -> None:
"""heuristic_v2 has its own slot, while custom tiers and custom prompts share one customization
slot: one v2 plus EITHER customization fits, but a second customization of any form is refused."""
router = Router(
model_list=[
self._POOL,
self._router_row("v2-a", "id-a", "heuristic_v2"),
self._custom_tier_row("tiers-a", "id-t"),
],
auto_router_capability_limit=lambda: 1,
ignore_invalid_deployments=True,
)
assert sorted(router.complexity_routers) == ["tiers-a", "v2-a"]
assert router.auto_router_capability_violation(HEURISTIC_V2_CAPABILITY) is not None
assert router.auto_router_capability_violation(CUSTOMIZATION_CAPABILITY) is not None
assert router.upsert_deployment(Deployment(**self._custom_tier_row("tiers-b", "id-t2"))) is None
assert router.upsert_deployment(Deployment(**self._custom_prompt_row("prompt-b", "id-p2"))) is None
assert router.upsert_deployment(Deployment(**self._router_row("v2-b", "id-b", "heuristic_v2"))) is None
assert sorted(router.complexity_routers) == ["tiers-a", "v2-a"]
@staticmethod
def _operator_prompt_row(model_name: str, model_id: str, field: str) -> dict[str, object]:
return {
"model_name": model_name,
"litellm_params": {
"model": "auto_router/complexity_router",
"complexity_router_default_model": "gpt-4o-mini",
"complexity_router_config": {
"classifier_type": "llm",
"classifier_llm_config": {"model": "gpt-4o-mini"},
"tiers": {"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o"},
field: '- "reset my password" -> SIMPLE',
},
},
"model_info": {"id": model_id},
}
@pytest.mark.parametrize("field", ["classification_prompt", "classification_examples"])
def test_operator_written_prompt_sections_claim_the_customization_slot(self, field: str) -> None:
"""The dashboard prompt editor writes opening instructions and calibration examples as their own
fields on a BUILT-IN tier router, so each must claim the slot on its own."""
with pytest.raises(ValueError, match="operator-written classifier prompt"):
Router(
model_list=[
self._POOL,
self._operator_prompt_row("prompt-a", "id-a", field),
self._operator_prompt_row("prompt-b", "id-b", field),
],
auto_router_capability_limit=lambda: 1,
)
@pytest.mark.parametrize("field", ["classification_prompt", "classification_examples"])
def test_an_operator_prompt_section_claims_the_slot_held_by_custom_tiers(self, field: str) -> None:
"""Switching the FORM of customization cannot buy a second unlicensed router."""
with pytest.raises(ValueError, match="operator-written classifier prompt"):
Router(
model_list=[
self._POOL,
self._custom_tier_row("tiers-a", "id-a"),
self._operator_prompt_row("prompt-b", "id-b", field),
],
auto_router_capability_limit=lambda: 1,
)
def test_a_custom_prompt_claims_the_slot_held_by_custom_tiers(self) -> None:
"""The customization ceiling is shared: changing its form cannot get a second unlicensed router."""
with pytest.raises(ValueError, match="operator-written classifier prompt"):
Router(
model_list=[
self._POOL,
self._custom_tier_row("tiers-a", "id-a"),
self._custom_prompt_row("prompt-b", "id-b"),
],
auto_router_capability_limit=lambda: 1,
)
def test_renaming_built_in_tiers_is_not_a_custom_tier_set(self) -> None:
"""tier_labels renames the built-in ladder without defining one, so it stays ungated: two such
routers register under a ceiling of one."""
def labeled(model_name: str, model_id: str) -> dict[str, object]:
row = self._router_row(model_name, model_id, "heuristic")
row["litellm_params"]["complexity_router_config"]["tier_labels"] = {"SIMPLE": "Cheap", "MEDIUM": "Standard"}
return row
router = Router(
model_list=[self._POOL, labeled("labels-a", "id-a"), labeled("labels-b", "id-b")],
auto_router_capability_limit=lambda: 1,
)
assert sorted(router.complexity_routers) == ["labels-a", "labels-b"]
def test_hybrid_initialization_waits_for_later_pool_deployments(self):
router = Router(
model_list=[

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

@ -5,9 +5,11 @@ 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,
GATED_AUTO_ROUTER_CAPABILITIES,
capability_limit_violation,
claimed_capability,
count_capability_routers,
gated_capability_of,
strategy_router_dependencies,
validate_complexity_router_config_placement,
validate_complexity_router_config_write,
@ -376,38 +378,122 @@ def test_placement_is_scoped_to_complexity_router_deployments(model, present_fie
assert carries_complexity_router_settings(model, present_fields) is scoped
_HV2_CONFIG: Mapping[str, object] = {"classifier_type": "heuristic_v2"}
_CUSTOM_TIER_CONFIG: Mapping[str, object] = {
"classifier_type": "llm",
"tier_definitions": [{"name": "routine", "description": "easy"}, {"name": "hard", "description": "hard"}],
}
_CUSTOM_PROMPT_CONFIG: Mapping[str, object] = {
"classifier_type": "llm",
"classifier_llm_config": {"model": "gpt-4o-mini", "system_prompt": "judge it my way"},
}
@pytest.mark.parametrize(
"litellm_params,expected",
"config,expected_key",
[
({"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),
(_CUSTOM_PROMPT_CONFIG, "tier_or_classifier_prompt"),
({"classifier_type": "llm", "classifier_llm_config": {"model": "m"}, "classification_prompt": "grade it"}, "tier_or_classifier_prompt"),
({"classifier_type": "llm", "classifier_llm_config": {"model": "m"}, "classification_examples": '- "x" -> SIMPLE'}, "tier_or_classifier_prompt"),
({"classifier_type": "hybrid", "classification_examples": "- y -> MEDIUM"}, "tier_or_classifier_prompt"),
({"classifier_type": "llm", "classifier_llm_config": {"model": "m"}, "classification_prompt": None, "classification_examples": None}, None),
({"classifier_type": "heuristic", "classification_examples": "- x -> SIMPLE"}, None),
({"classifier_type": "hybrid", "classifier_llm_config": {"system_prompt": "p"}}, "tier_or_classifier_prompt"),
({"classifier_type": "heuristic_first", "classifier_llm_config": {"system_prompt": "p"}}, "tier_or_classifier_prompt"),
({"classifier_type": "llm", "classifier_llm_config": {"model": "m", "classification_rubric": "chat"}}, None),
({"classifier_type": "llm", "classifier_llm_config": {"model": "m"}}, None),
({"classifier_type": "llm", "classifier_llm_config": {"model": "m", "system_prompt": None}}, None),
({"classifier_type": "heuristic", "classifier_llm_config": {"system_prompt": "p"}}, None),
({"classifier_type": "heuristic_v2", "classifier_llm_config": {"system_prompt": "p"}}, "heuristic_v2"),
({"classifier_type": "llm", "classifier_llm_config": "not a mapping"}, None),
],
)
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_custom_classifier_prompt_capability(config: Mapping[str, object], expected_key: str | None) -> None:
"""Every operator-written part of the classifier prompt claims the customization slot: a whole
replacement system_prompt, replacement opening instructions (classification_prompt), or replacement
calibration examples (classification_examples).
A shipped rubric preset stays free, and the heuristic scorers never read system_prompt, so a
value sitting on one is inert and claims nothing (heuristic_v2 still claims its own capability).
"""
claimed = claimed_capability(config)
assert (None if claimed is None else claimed.key) == expected_key
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"}}
@pytest.mark.parametrize(
"model,expected",
[
("auto_router/complexity_router", True),
("auto_router/complexity_router-eu", True),
("auto_router/semantic_router", False),
("auto_router/adaptive_router", False),
("auto_router/quality_router", False),
("openai/gpt-4o", False),
(None, False),
],
)
def test_is_complexity_router_model(model: str | None, expected: bool) -> None:
from litellm.router_utils.auto_router_model_naming import is_complexity_router_model
assert is_complexity_router_model(model) is expected
@pytest.mark.parametrize(
"litellm_params,expected_key",
[
({"model": "auto_router/complexity_router", "complexity_router_config": _HV2_CONFIG}, "heuristic_v2"),
({"model": "auto_router/complexity_router-eu", "complexity_router_config": _HV2_CONFIG}, "heuristic_v2"),
({"model": "auto_router/complexity_router", "complexity_router_config": _CUSTOM_TIER_CONFIG}, "tier_or_classifier_prompt"),
({"model": "auto_router/complexity_router-eu", "complexity_router_config": _CUSTOM_TIER_CONFIG}, "tier_or_classifier_prompt"),
({"model": "auto_router/complexity_router", "complexity_router_config": {"classifier_type": "heuristic"}}, None),
({"model": "auto_router/complexity_router", "complexity_router_config": {"tiers": {"SIMPLE": "a"}}}, None),
({"model": "auto_router/complexity_router", "complexity_router_config": {"tier_definitions": None}}, None),
({"model": "auto_router/complexity_router", "complexity_router_config": {"tier_labels": {"SIMPLE": "Cheap"}}}, None),
({"model": "auto_router/complexity_router"}, None),
({"model": "auto_router/quality_router", "complexity_router_config": _HV2_CONFIG}, None),
({"model": "auto_router/quality_router", "complexity_router_config": _CUSTOM_TIER_CONFIG}, None),
({"model": "openai/gpt-4o", "complexity_router_config": _HV2_CONFIG}, None),
({"model": "openai/gpt-4o", "complexity_router_config": _CUSTOM_TIER_CONFIG}, None),
({"model": "auto_router/complexity_router", "complexity_router_config": "heuristic_v2"}, None),
({}, None),
],
)
def test_gated_capability_of(litellm_params: Mapping[str, object], expected_key: str | None) -> None:
"""Only a complexity router claiming a licensed capability counts toward that capability's limit.
Renaming the built-in tiers through tier_labels is not a custom tier set, so it stays ungated.
"""
capability = gated_capability_of(litellm_params)
assert (None if capability is None else capability.key) == expected_key
@pytest.mark.parametrize("capability", GATED_AUTO_ROUTER_CAPABILITIES, ids=lambda c: c.key)
def test_count_capability_routers_counts_only_its_own_capability(capability) -> None:
"""Each capability has its own ceiling, so a router claiming the sibling capability never counts,
while a custom tier set and a custom classifier prompt count into the SAME customization slot."""
def row(name: str, config: Mapping[str, object] | None) -> Mapping[str, object]:
params = {"model": "auto_router/complexity_router"} | ({} if config is None else {"complexity_router_config": config})
return {"model_name": name, "litellm_params": params}
by_key = {
"heuristic_v2": (_HV2_CONFIG, _HV2_CONFIG),
"tier_or_classifier_prompt": (_CUSTOM_TIER_CONFIG, _CUSTOM_PROMPT_CONFIG),
}
mine_first, mine_second = by_key[capability.key]
theirs = next(configs[0] for key, configs in by_key.items() if key != capability.key)
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"},
row("a", mine_first),
row("b", theirs),
{"model_name": "c", "litellm_params": {"model": "openai/gpt-4o"}},
row("d", mine_second),
{"model_name": "e"},
{"model_name": "f", "litellm_params": "not a mapping"},
]
assert count_heuristic_v2_routers(rows) == 2
assert count_heuristic_v2_routers(()) == 0
assert count_capability_routers(rows, capability=capability) == 2
assert count_capability_routers((), capability=capability) == 0
@pytest.mark.parametrize("capability", GATED_AUTO_ROUTER_CAPABILITIES, ids=lambda c: c.key)
@pytest.mark.parametrize(
"held,limit,violates",
[
@ -419,10 +505,42 @@ def test_count_heuristic_v2_routers_reads_model_list_rows_and_ignores_malformed_
(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)
def test_capability_limit_violation(held: int, limit: int | None, violates: bool, capability) -> None:
violation = capability_limit_violation(capability=capability, 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 capability.subject in violation
assert capability.remedy in violation
assert "license" not in violation
def test_every_gated_capability_has_a_distinct_predicate_and_sql_spelling() -> None:
"""The in-process and SQL halves of a capability must stay paired, and no two capabilities may collide."""
keys = tuple(capability.key for capability in GATED_AUTO_ROUTER_CAPABILITIES)
assert len(set(keys)) == len(keys)
for capability in GATED_AUTO_ROUTER_CAPABILITIES:
assert "{config}" in capability.sql_config_predicate
assert capability.uses is not None
@pytest.mark.parametrize(
"config",
[
_HV2_CONFIG,
_CUSTOM_TIER_CONFIG,
_CUSTOM_PROMPT_CONFIG,
{"classifier_type": "heuristic"},
{"classifier_type": "heuristic_v2", "classifier_llm_config": {"system_prompt": "p"}},
{"classifier_type": "llm", "classifier_llm_config": {"model": "m", "system_prompt": "p"}, "tier_labels": {"SIMPLE": "Cheap"}},
],
)
def test_capabilities_are_mutually_exclusive_on_one_config(config: Mapping[str, object]) -> None:
"""No config claims two capabilities, which is what lets one lock and one count serve them all.
The config validator is what makes this true and is pinned separately in test_complexity_router:
tier_definitions rejects every heuristic classifier_type and rejects the classifier system_prompt,
and system_prompt only counts for the classifier types heuristic_v2 is not one of.
"""
assert sum(1 for capability in GATED_AUTO_ROUTER_CAPABILITIES if capability.uses(config)) <= 1