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

litellm/constants.py

@@ -40,7 +40,7 @@
         "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))

litellm/proxy/auth/litellm_license.py

@@ -17,7 +17,7 @@
 
 
 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 @@ def is_team_count_over_limit(self, team_count: int) -> bool:
             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

litellm/proxy/management_endpoints/model_management_endpoints.py

@@ -50,7 +50,7 @@
     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 @@
     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
 
 
+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
+
+
+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 _heuristic_v2_slot(
-    prisma_client: PrismaClient, *, effective_config: object, model_id: str | None
+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 heuristic_v2 slot is settled.
+    """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 on classifier_type heuristic_v2 under a limited license runs
+    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 write_row(update_data: PrismaCompatibleUpdateDBModel) -> _ProxyModelRo
             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(

litellm/proxy/proxy_server.py

@@ -118,10 +118,11 @@
     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 @@ def generate_feedback_box():
 )
 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 @@ async def load_config(self, router: litellm.Router | None, config_file_path: str
         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 @@ async def load_config(self, router: litellm.Router | None, config_file_path: str
             ),
             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 @@ async def _update_llm_router(
                         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: