From 0a6689e6f8d5914e303252542ea2799b5001bc2b Mon Sep 17 00:00:00 2001 From: moe-berri Date: Tue, 1 Sep 2026 15:31:17 -0700 Subject: [PATCH] feat(proxy): manage capability router configs --- litellm/proxy/_types.py | 1 + .../auto_router_endpoints.py | 78 +++++++++++++++---- .../model_management_endpoints.py | 7 ++ .../auto_router_endpoints.py | 24 +++++- .../proxy/auth/test_route_checks.py | 19 +++-- .../test_capability_router_endpoints.py | 60 ++++++++++++++ 6 files changed, 164 insertions(+), 25 deletions(-) create mode 100644 tests/test_litellm/proxy/management_endpoints/test_capability_router_endpoints.py diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 18714256a8f..d7a242cf5f8 100644 --- a/litellm/proxy/_types.py +++ b/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, diff --git a/litellm/proxy/management_endpoints/auto_router_endpoints.py b/litellm/proxy/management_endpoints/auto_router_endpoints.py index 21e652114bc..5d0395ea7f1 100644 --- a/litellm/proxy/management_endpoints/auto_router_endpoints.py +++ b/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"], diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index ca66640bf46..adce05f02c5 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/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) diff --git a/litellm/types/management_endpoints/auto_router_endpoints.py b/litellm/types/management_endpoints/auto_router_endpoints.py index 88869a1edfb..137f136e5f1 100644 --- a/litellm/types/management_endpoints/auto_router_endpoints.py +++ b/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( diff --git a/tests/test_litellm/proxy/auth/test_route_checks.py b/tests/test_litellm/proxy/auth/test_route_checks.py index 71ccef620e5..158fe908f6f 100644 --- a/tests/test_litellm/proxy/auth/test_route_checks.py +++ b/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 diff --git a/tests/test_litellm/proxy/management_endpoints/test_capability_router_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_capability_router_endpoints.py new file mode 100644 index 00000000000..40bc88351af --- /dev/null +++ b/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