feat(proxy): manage capability router configs

litellm/proxy/_types.py

@@ -860,6 +860,7 @@ class LiteLLMRoutes(enum.Enum):
         # proxy admin, or team admin naming their own team via team_id
         "/auto_router/test_routing",
         "/auto_router/validate_complexity_router_config",
+        "/auto_router/validate_capability_router_config",
         # Agent registry - reads are role-scoped and writes are proxy-admin-gated
         # inside agent_endpoints/endpoints.py
         *agent_management_routes,

litellm/proxy/management_endpoints/auto_router_endpoints.py

@@ -39,6 +39,7 @@ from litellm.proxy.litellm_pre_call_utils import (
 )
 from litellm.repositories.base_repository import SupportsModelDump
 from litellm.repositories.team_repository import TeamRepository
+from litellm.router_strategy.capability_router import CapabilityRouter
 from litellm.router_strategy.complexity_router import ComplexityRouter
 from litellm.router_utils.auto_router_model_naming import (
     StrategyRouterDependencyRole,
@@ -54,6 +55,8 @@ from litellm.types.management_endpoints.auto_router_endpoints import (
     AutoRouterCacheStats,
     AutoRouterRoutingTestRequest,
     AutoRouterRoutingTestResponse,
+    CapabilityRouterConfigValidationRequest,
+    CapabilityRouterConfigValidationResponse,
     ComplexityRouterConfigValidationRequest,
     ComplexityRouterConfigValidationResponse,
     RequestComplexityRouterConfig,
@@ -277,7 +280,19 @@ async def _authorize_models_this_test_can_call(
     the key's own budget is not checked either. Test Connection gets both for free by routing
     its calls through the proxy. Team and member budgets are already enforced on every route.
     """
-    models: Final = _models_this_test_can_call(config)
+    await _authorize_model_names_this_test_can_call(
+        models=_models_this_test_can_call(config),
+        user_api_key_dict=user_api_key_dict,
+        llm_router=llm_router,
+    )
+
+
+async def _authorize_model_names_this_test_can_call(
+    models: Sequence[str],
+    user_api_key_dict: UserAPIKeyAuth,
+    llm_router: "Router",
+) -> None:
+    """Apply model-access and key-budget checks to internal dry-run calls."""
     if not models:
         return
 
@@ -334,6 +349,27 @@ async def validate_complexity_router_config(
     return ComplexityRouterConfigValidationResponse(valid=error is None, error=error)
 
 
+@router.post(
+    "/auto_router/validate_capability_router_config",
+    tags=["model management"],
+    dependencies=[Depends(user_api_key_auth)],
+    response_model=CapabilityRouterConfigValidationResponse,
+    status_code=status.HTTP_200_OK,
+)
+async def validate_capability_router_config(
+    data: CapabilityRouterConfigValidationRequest,
+    user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
+) -> CapabilityRouterConfigValidationResponse:
+    """Validate a capability-router config without saving it."""
+    await _authorize_router_dry_run(user_api_key_dict=user_api_key_dict, team_id=data.team_id)
+    from litellm.router_utils.auto_router_model_naming import (
+        validate_capability_router_config_write,
+    )
+
+    error: Final = validate_capability_router_config_write(data.capability_router_config)
+    return CapabilityRouterConfigValidationResponse(valid=error is None, error=error)
+
+
 @router.post(
     "/auto_router/test_routing",
     tags=["model management"],  # mutable-ok: fastapi's decorator signature types tags as a list
@@ -397,19 +433,31 @@ async def preview_auto_router_routing(
             },
         )
 
-    await _authorize_models_this_test_can_call(
-        config=data.complexity_router_config,
-        user_api_key_dict=user_api_key_dict,
-        llm_router=llm_router,
-    )
-
-    complexity_router: Final = ComplexityRouter(
-        model_name=data.router_name,
-        litellm_router_instance=llm_router,
-        complexity_router_config=data.complexity_router_config.model_dump(exclude_none=True),
-        default_model=data.default_model,
-        derive_savings_baseline=False,
-    )
+    if data.capability_router_config is not None:
+        await _authorize_model_names_this_test_can_call(
+            models=(data.capability_router_config.classifier.model,),
+            user_api_key_dict=user_api_key_dict,
+            llm_router=llm_router,
+        )
+        strategy = CapabilityRouter(
+            model_name=data.router_name,
+            litellm_router_instance=llm_router,
+            capability_router_config=data.capability_router_config.model_dump(exclude_none=True),
+        )
+    else:
+        assert data.complexity_router_config is not None
+        await _authorize_models_this_test_can_call(
+            config=data.complexity_router_config,
+            user_api_key_dict=user_api_key_dict,
+            llm_router=llm_router,
+        )
+        strategy = ComplexityRouter(
+            model_name=data.router_name,
+            litellm_router_instance=llm_router,
+            complexity_router_config=data.complexity_router_config.model_dump(exclude_none=True),
+            default_model=data.default_model,
+            derive_savings_baseline=False,
+        )
 
     request_kwargs: Final = LiteLLMProxyRequestSetup.add_user_api_key_auth_to_request_metadata(
         data={  # mutable-ok: the request-metadata helper takes and returns request kwargs as a dict
@@ -423,7 +471,7 @@ async def preview_auto_router_routing(
     refresh_proxy_server_request_body_snapshot(request_kwargs)
 
     try:
-        hook_response: Final = await complexity_router.async_pre_routing_hook(
+        hook_response: Final = await strategy.async_pre_routing_hook(
             model=data.router_name,
             request_kwargs=request_kwargs,
             messages=request_kwargs["messages"],

litellm/proxy/management_endpoints/model_management_endpoints.py

@@ -92,6 +92,7 @@ from litellm.router_strategy.complexity_router import (
 from litellm.router_utils.auto_router_model_naming import (
     STRATEGY_ROUTER_PARAM_FIELDS,
     carries_complexity_router_settings,
+    validate_capability_router_config_write,
     validate_complexity_router_config_placement,
     validate_complexity_router_config_write,
     validate_strategy_router_model_write,
@@ -229,6 +230,11 @@ def _strategy_router_write_violation(
     """
     if incoming_params is None:
         return None
+    capability_violation: Final = validate_capability_router_config_write(
+        capability_router_config=incoming_params.capability_router_config
+    )
+    if capability_violation is not None:
+        return capability_violation
     config_violation: Final = validate_complexity_router_config_write(
         complexity_router_config=incoming_params.complexity_router_config
     )
@@ -2608,6 +2614,7 @@ async def clear_cache() -> ReconcileOutcome:
             # on reload and abort it.
             for model_name in db_router_names:
                 llm_router.auto_routers.pop(model_name, None)
+                llm_router.capability_routers.pop(model_name, None)
                 llm_router.complexity_routers.pop(model_name, None)
                 llm_router.adaptive_routers.pop(model_name, None)
                 llm_router.quality_routers.pop(model_name, None)

litellm/types/management_endpoints/auto_router_endpoints.py

@@ -9,6 +9,7 @@ from typing import Final, Literal, TypeAlias
 
 from pydantic import BaseModel, Field, computed_field, field_validator, model_validator
 
+from litellm.router_strategy.capability_router.config import CapabilityRouterConfig
 from litellm.router_strategy.complexity_router.config import ComplexityRouterConfig
 from litellm.types.utils import StandardLoggingRoutingDecision
 
@@ -44,6 +45,18 @@ class ComplexityRouterConfigValidationResponse(BaseModel):
     error: str | None = None
 
 
+class CapabilityRouterConfigValidationRequest(BaseModel):
+    """A capability-router config to validate without saving."""
+
+    capability_router_config: Mapping[str, object]
+    team_id: str | None = None
+
+
+class CapabilityRouterConfigValidationResponse(BaseModel):
+    valid: bool
+    error: str | None = None
+
+
 class AutoRouterRoutingTestRequest(BaseModel):
     """A single request to classify against a complexity-router config that need not be saved yet.
 
@@ -69,8 +82,13 @@ class AutoRouterRoutingTestRequest(BaseModel):
         default=None,
         description="The tool definitions the request advertises, which decide whether the plan-mode floor applies",
     )
-    complexity_router_config: RequestComplexityRouterConfig = Field(
-        description="The complexity router config to route against, in the shape /model/new accepts",
+    complexity_router_config: RequestComplexityRouterConfig | None = Field(
+        default=None,
+        description="The complexity router config to route against",
+    )
+    capability_router_config: CapabilityRouterConfig | None = Field(
+        default=None,
+        description="The capability router config to route against",
     )
     default_model: str | None = Field(
         default=None,
@@ -114,6 +132,8 @@ class AutoRouterRoutingTestRequest(BaseModel):
             raise ValueError("messages must not be empty")
         if (self.prompt is None) == (self.messages is None):
             raise ValueError("provide exactly one of prompt or messages")
+        if (self.complexity_router_config is None) == (self.capability_router_config is None):
+            raise ValueError("provide exactly one router config")
         if self.messages is not None:
             return self
         return self.model_copy(

tests/test_litellm/proxy/auth/test_route_checks.py

@@ -2,7 +2,6 @@ import os
 from datetime import datetime
 from unittest.mock import MagicMock, patch
 
-
 import pytest
 from fastapi import HTTPException, Request
 
@@ -803,7 +802,7 @@ def test_google_routes_with_dynamic_model_names_accessible_to_internal_users():
         # If no exception is raised, the test passes
     except Exception as e:
         pytest.fail(
-            f"Internal user should be able to access Google generateContent route. Got error: {str(e)}"
+            f"Internal user should be able to access Google generateContent route. Got error: {e!s}"
         )
 
 
@@ -1703,7 +1702,7 @@ def test_videos_route_accessible_to_internal_users():
         # If no exception is raised, the test passes
     except Exception as e:
         pytest.fail(
-            f"Internal user should be able to access /v1/videos route. Got error: {str(e)}"
+            f"Internal user should be able to access /v1/videos route. Got error: {e!s}"
         )
 
 
@@ -1803,7 +1802,7 @@ def test_proxy_admin_viewer_can_access_global_spend_tags():
         # If no exception is raised, the test passes
     except Exception as e:
         pytest.fail(
-            f"proxy_admin_viewer should be able to access /global/spend/tags route. Got error: {str(e)}"
+            f"proxy_admin_viewer should be able to access /global/spend/tags route. Got error: {e!s}"
         )
 
 
@@ -1964,7 +1963,7 @@ def test_proxy_admin_viewer_can_access_audit_logs(route):
         )
     except Exception as e:
         pytest.fail(
-            f"proxy_admin_viewer should be able to access {route} route. Got error: {str(e)}"
+            f"proxy_admin_viewer should be able to access {route} route. Got error: {e!s}"
         )
 
 
@@ -2029,7 +2028,7 @@ def test_proxy_admin_viewer_can_access_logs_page_endpoints(route):
         )
     except Exception as e:
         pytest.fail(
-            f"proxy_admin_viewer should be able to access {route}. Got error: {str(e)}"
+            f"proxy_admin_viewer should be able to access {route}. Got error: {e!s}"
         )
 
 
@@ -2141,7 +2140,7 @@ def test_proxy_admin_viewer_can_access_settings_read_endpoints(route):
         )
     except Exception as e:
         pytest.fail(
-            f"proxy_admin_viewer should be able to access {route}. Got error: {str(e)}"
+            f"proxy_admin_viewer should be able to access {route}. Got error: {e!s}"
         )
 
 
@@ -3407,7 +3406,11 @@ def test_organization_daily_activity_not_granted_by_org_admin_request_data_branc
 )
 @pytest.mark.parametrize(
     "dry_run_route",
-    ["/auto_router/test_routing", "/auto_router/validate_complexity_router_config"],
+    [
+        "/auto_router/test_routing",
+        "/auto_router/validate_complexity_router_config",
+        "/auto_router/validate_capability_router_config",
+    ],
 )
 def test_auto_router_dry_runs_share_model_new_audience(user_role, dry_run_route):
     """The dry runs serve whoever can draft a save on /model/new, no one else: a role

tests/test_litellm/proxy/management_endpoints/test_capability_router_endpoints.py

@@ -0,0 +1,60 @@
+import pytest
+from pydantic import ValidationError
+
+from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
+from litellm.proxy.management_endpoints.auto_router_endpoints import validate_capability_router_config
+from litellm.types.management_endpoints.auto_router_endpoints import (
+    AutoRouterRoutingTestRequest,
+    CapabilityRouterConfigValidationRequest,
+)
+
+
+def config() -> dict:
+    return {
+        "candidates": [
+            {"model": "small", "description": "Reliable for short extraction tasks"},
+            {"model": "frontier", "description": "Reliable for ambiguous multi-step tasks"},
+        ],
+        "classifier": {"model": "classifier"},
+        "probability_threshold": 0.7,
+        "fallback_model": "frontier",
+    }
+
+
+def test_routing_preview_accepts_exactly_one_router_config() -> None:
+    request = AutoRouterRoutingTestRequest.model_validate(
+        {"prompt": "Extract the invoice number", "capability_router_config": config()}
+    )
+    assert request.capability_router_config is not None
+    assert request.complexity_router_config is None
+
+    with pytest.raises(ValidationError, match="exactly one router config"):
+        AutoRouterRoutingTestRequest.model_validate(
+            {
+                "prompt": "Extract the invoice number",
+                "capability_router_config": config(),
+                "complexity_router_config": {
+                    "tiers": {"SIMPLE": "small"},
+                    "classifier_type": "heuristic",
+                },
+            }
+        )
+
+
+@pytest.mark.asyncio
+async def test_validation_endpoint_uses_runtime_config_contract() -> None:
+    admin = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, api_key="admin")
+    accepted = await validate_capability_router_config(
+        CapabilityRouterConfigValidationRequest(capability_router_config=config()),
+        admin,
+    )
+    rejected = await validate_capability_router_config(
+        CapabilityRouterConfigValidationRequest(
+            capability_router_config={**config(), "fallback_model": "missing"}
+        ),
+        admin,
+    )
+
+    assert accepted.valid is True
+    assert rejected.valid is False
+    assert rejected.error is not None and "fallback_model" in rejected.error