From 97c54e278e1da08f3c770a554c67fb2a47eb1424 Mon Sep 17 00:00:00 2001 From: Moe Khalil Date: Sun, 20 Sep 2026 00:54:17 +0000 Subject: [PATCH] fix(auto-router): resolve saved JEV probes on the server Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../auto_router_endpoints.py | 68 ++++++++++---- .../auto_router_endpoints.py | 5 + .../test_auto_router_endpoints.py | 93 ++++++++++++++++++- .../JevConnectionTest.integration.test.tsx | 13 ++- ...d_auto_router_routing_test_request.test.ts | 30 ++++-- .../build_auto_router_routing_test_request.ts | 9 +- .../src/components/model_info_view.tsx | 3 +- .../src/components/networking.tsx | 1 + ui/litellm-dashboard/src/lib/http/schema.d.ts | 5 + 9 files changed, 187 insertions(+), 40 deletions(-) diff --git a/litellm/proxy/management_endpoints/auto_router_endpoints.py b/litellm/proxy/management_endpoints/auto_router_endpoints.py index 19d6d9b4e42..07dee3edf15 100644 --- a/litellm/proxy/management_endpoints/auto_router_endpoints.py +++ b/litellm/proxy/management_endpoints/auto_router_endpoints.py @@ -384,6 +384,40 @@ async def validate_complexity_router_config( return ComplexityRouterConfigValidationResponse(valid=error is None, error=error) +async def _resolve_saved_routing_test( + data: AutoRouterRoutingTestRequest, + user_api_key_dict: UserAPIKeyAuth, + llm_router: "Router", +) -> AutoRouterRoutingTestRequest: + if data.saved_model_id is None: + return data + deployment: Final = llm_router.get_deployment(data.saved_model_id) + if deployment is None or deployment.model_info.blocked: + raise HTTPException(status_code=404, detail="Saved auto router is unavailable") + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN and deployment.model_info.team_id != data.team_id: + raise HTTPException(status_code=403, detail="Saved auto router belongs to a different team") + await can_key_call_resolved_model( + model=deployment.model_info.team_public_model_name or deployment.model_name, + llm_model_list=llm_router.model_list, + valid_token=user_api_key_dict, + llm_router=llm_router, + ) + params: Final = deployment.litellm_params + if classify_strategy_router_model(params.model or "") != "complexity" or params.complexity_router_config is None: + raise HTTPException(status_code=400, detail="Saved deployment is not a complexity auto router") + return data.model_copy( + update=MappingProxyType( + { + "complexity_router_config": RequestComplexityRouterConfig.model_validate( + params.complexity_router_config + ), + "default_model": params.complexity_router_default_model, + "router_name": deployment.model_name, + } + ) + ) + + @router.post( "/auto_router/test_routing", tags=["model management"], # mutable-ok: fastapi's decorator signature types tags as a list @@ -439,10 +473,18 @@ async def preview_auto_router_routing( from litellm.proxy.utils import get_available_models_for_user member_team: Final = await _authorize_router_dry_run(user_api_key_dict=user_api_key_dict, team_id=data.team_id) + if llm_router is None: + raise HTTPException( + status_code=500, + detail={ # mutable-ok: HTTPException detail must be a plain mapping + "error": CommonProxyErrors.no_llm_router.value + }, + ) + resolved: Final = await _resolve_saved_routing_test(data, user_api_key_dict, llm_router) actor: Final = ( await _authorize_member_dry_run_config( - config=data.complexity_router_config.model_dump(exclude_none=True), - default_model=data.default_model, + config=resolved.complexity_router_config.model_dump(exclude_none=True), + default_model=resolved.default_model, user_api_key_dict=user_api_key_dict, team=member_team, ) @@ -450,12 +492,12 @@ async def preview_auto_router_routing( else user_api_key_dict ) request_data: Final[dict[str, object]] = { # mutable-ok: auth and routing enrich this request in place - **data.wire_body(), + **resolved.wire_body(), "metadata": {}, # mutable-ok: centralized auth and identity stamping share this metadata bucket "proxy_server_request": {"body": None}, # mutable-ok: the snapshot owner fills this body in place } - if member_team is not None and _models_this_test_can_call(data.complexity_router_config): + if member_team is not None and _models_this_test_can_call(resolved.complexity_router_config): from litellm.proxy.auth.user_api_key_auth import ( _run_centralized_common_checks, # pyright: ignore[reportPrivateUsage] # reuse the serving admission policy ) @@ -467,25 +509,17 @@ async def preview_auto_router_routing( route="/auto_router/test_routing", ) - if llm_router is None: - raise HTTPException( - status_code=500, - detail={ # mutable-ok: HTTPException detail must be a plain mapping - "error": CommonProxyErrors.no_llm_router.value - }, - ) - await _authorize_models_this_test_can_call( - config=data.complexity_router_config, + config=resolved.complexity_router_config, user_api_key_dict=actor, llm_router=llm_router, ) complexity_router: Final = ComplexityRouter( - model_name=data.router_name, + model_name=resolved.router_name, litellm_router_instance=llm_router, - complexity_router_config=data.complexity_router_config.model_dump(exclude_none=True), - default_model=data.default_model, + complexity_router_config=resolved.complexity_router_config.model_dump(exclude_none=True), + default_model=resolved.default_model, derive_savings_baseline=False, ) @@ -498,7 +532,7 @@ async def preview_auto_router_routing( try: hook_response: Final = await complexity_router.async_pre_routing_hook( - model=data.router_name, + model=resolved.router_name, request_kwargs=request_kwargs, messages=request_kwargs["messages"], ) diff --git a/litellm/types/management_endpoints/auto_router_endpoints.py b/litellm/types/management_endpoints/auto_router_endpoints.py index fd2202a1156..93ea925bd9e 100644 --- a/litellm/types/management_endpoints/auto_router_endpoints.py +++ b/litellm/types/management_endpoints/auto_router_endpoints.py @@ -72,6 +72,11 @@ class AutoRouterRoutingTestRequest(BaseModel): complexity_router_config: RequestComplexityRouterConfig = Field( description="The complexity router config to route against, in the shape /model/new accepts", ) + saved_model_id: str | None = Field( + default=None, + min_length=1, + description="Test this saved deployment's server-side configuration instead of the supplied config and default model", + ) default_model: str | None = Field( default=None, description="Model to route to when no tier resolves, i.e. complexity_router_default_model", diff --git a/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py index f9b618234b6..9235a00bda6 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py @@ -7,13 +7,13 @@ from pathlib import Path from typing import Final import httpx -import litellm.llms.custom_httpx.http_handler as http_handler -import litellm.router_strategy.complexity_router.complexity_router as complexity_module import pytest import respx from fastapi import HTTPException, Request from pydantic import ValidationError +import litellm.llms.custom_httpx.http_handler as http_handler +import litellm.router_strategy.complexity_router.complexity_router as complexity_module from litellm.proxy import proxy_server from litellm.proxy._types import ( LitellmUserRoles, @@ -29,6 +29,7 @@ from litellm.types.management_endpoints.auto_router_endpoints import ( AutoRouterBenchmarksResponse, AutoRouterRoutingTestRequest, ) +from litellm.types.router import Deployment from litellm.types.utils import Choices, Message, ModelResponse ROUTING_HTTP_REQUEST: Final = Request( @@ -2382,6 +2383,94 @@ async def test_jev_test_routing_authorizes_paid_evaluation_before_contacting_typ await handler.client.aclose() +@pytest.mark.asyncio +@pytest.mark.parametrize("case", ["allowed", "missing", "blocked", "key", "budget", "team", "not-router"]) +async def test_saved_jev_probe_uses_authorized_server_configuration(monkeypatch: pytest.MonkeyPatch, case: str) -> None: + router: Final = RecordingRouter("SIMPLE") + stored_key: Final = "synthetic-server-jev-key" + stored_config: Final = { + "classifier_type": "jev", + "tiers": TIERS, + "jev_classifier_config": {"api_key": stored_key, "api_base": "https://saved-jev.test"}, + } + router.add_deployment( + Deployment.model_validate( + { + "model_name": "saved-jev", + "litellm_params": { + "model": "openai/gpt-4o-mini" if case == "not-router" else "auto_router/complexity_router", + "complexity_router_config": stored_config, + }, + "model_info": { + "id": "saved-jev-id", + "blocked": case == "blocked", + "team_id": "owner-team" if case == "team" else None, + }, + } + ) + ) + monkeypatch.setattr(proxy_server, "llm_router", router) + actor: Final = ( + _configure_member_preview(monkeypatch) + if case == "team" + else UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-probe", + user_id="admin", + models=["typesafe/jev-latest"] if case == "key" else ["saved-jev", "typesafe/jev-latest"], + max_budget=1, + spend=1 if case == "budget" else 0, + ) + ) + request: Final = _request_from( + { + "prompt": "what is 2+2", + "saved_model_id": "missing-id" if case == "missing" else "saved-jev-id", + "team_id": "member-preview-team" if case == "team" else None, + }, + classifier_type="jev", + jev_classifier_config={"api_key": "masked-key", "api_base": "https://browser-override.test"}, + ) + with respx.mock(assert_all_called=False) as http: + handler: Final = http_handler.AsyncHTTPHandler() + handler.client = httpx.AsyncClient(transport=httpx.MockTransport(http.async_handler)) + + def http_client(_provider: object) -> http_handler.AsyncHTTPHandler: + return handler + + monkeypatch.setattr(complexity_module, "get_async_httpx_client", http_client) + evaluation: Final = http.post("https://saved-jev.test/v1/systemone").mock( + return_value=httpx.Response( + 200, + json={ + "answers": { + "tier": {"type": "choice", "choice": "SIMPLE", "confidence": 1, "probabilities": {"SIMPLE": 1}} + } + }, + ) + ) + operation: Final = preview_auto_router_routing(request, actor, ROUTING_HTTP_REQUEST) + if case in ("missing", "blocked", "team", "not-router"): + with pytest.raises(HTTPException) as denied: + await operation + assert denied.value.status_code == {"missing": 404, "blocked": 404, "team": 403, "not-router": 400}[case] + elif case in ("key", "budget"): + with pytest.raises(ProxyException) as forbidden: + await operation + assert forbidden.value.type == ( + ProxyErrorTypes.key_model_access_denied if case == "key" else ProxyErrorTypes.budget_exceeded + ) + else: + result: Final = await operation + assert result.routing_decision["cause"] == "jev_classifier" + assert result.routed_model == "cheap-model" + assert evaluation.calls.last.request.headers["authorization"] == f"Bearer {stored_key}" + assert stored_key not in result.model_dump_json() + assert evaluation.call_count == (1 if case == "allowed" else 0) + assert router.recorded_calls == [] + await handler.client.aclose() + + @pytest.mark.asyncio async def test_list_shadow_eval_jobs_filters_to_jobs_containing_the_key(monkeypatch: pytest.MonkeyPatch): """The filter matches a key anywhere in a job's key set and still returns the whole diff --git a/ui/litellm-dashboard/src/components/add_model/JevConnectionTest.integration.test.tsx b/ui/litellm-dashboard/src/components/add_model/JevConnectionTest.integration.test.tsx index 2a00e8bb45e..72acda7622a 100644 --- a/ui/litellm-dashboard/src/components/add_model/JevConnectionTest.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/JevConnectionTest.integration.test.tsx @@ -45,7 +45,13 @@ const configParams: BuildComplexityRouterConfigParams = { returnRawModelName: false, }; const config = buildComplexityRouterConfig(configParams); -const request = buildSavedJevConnectionTestRequest(JSON.stringify(config), "fast", "my-router"); +const request = buildSavedJevConnectionTestRequest( + JSON.stringify({ + ...config, + jev_classifier_config: { api_key: "sk-masked****", api_base: "https://custom-jev.test" }, + }), + "saved-id", +); const targets = buildAutoRouterTestTargets({ tiers: Object.entries(config.tiers), semanticMatchingEnabled: false, @@ -95,9 +101,8 @@ describe("JEV network probes", () => { const routingCall = fetchMock.mock.calls.find(([url]) => String(url).endsWith("/auto_router/test_routing")); const expectedRequest = { prompt: JEV_CONNECTION_TEST_PROMPT, - complexity_router_config: config, - default_model: "fast", - router_name: "my-router", + complexity_router_config: { ...config, jev_classifier_config: undefined }, + saved_model_id: "saved-id", }; expect(JSON.parse(String(routingCall?.[1]?.body))).toEqual(expectedRequest); expect(fetchMock).toHaveBeenCalledTimes(5); diff --git a/ui/litellm-dashboard/src/components/add_model/build_auto_router_routing_test_request.test.ts b/ui/litellm-dashboard/src/components/add_model/build_auto_router_routing_test_request.test.ts index fba4ca47e00..174f93eae6c 100644 --- a/ui/litellm-dashboard/src/components/add_model/build_auto_router_routing_test_request.test.ts +++ b/ui/litellm-dashboard/src/components/add_model/build_auto_router_routing_test_request.test.ts @@ -20,6 +20,22 @@ const params = { }; describe("buildAutoRouterRoutingTestRequest", () => { + it("references the saved deployment without copying masked credentials or client overrides", () => { + const request = buildSavedJevConnectionTestRequest( + { + classifier_type: "jev", + tiers: CONFIG.tiers, + jev_classifier_config: { api_key: "sk-masked****", api_base: "https://custom-jev.test" }, + }, + "saved-id", + ); + const expectedRequest = { + prompt: JEV_CONNECTION_TEST_PROMPT, + complexity_router_config: { classifier_type: "jev", tiers: CONFIG.tiers }, + saved_model_id: "saved-id", + }; + expect(request).toEqual(expectedRequest); + }); it.each(["object", "json"])("probes saved JEV %s configuration with custom tiers and team context", (format) => { const config = { classifier_type: "jev", @@ -31,24 +47,18 @@ describe("buildAutoRouterRoutingTestRequest", () => { }; const expectedRequest = { prompt: JEV_CONNECTION_TEST_PROMPT, - complexity_router_config: config, - default_model: "strong", - router_name: "saved-router", + complexity_router_config: { ...config, jev_classifier_config: undefined }, + saved_model_id: "saved-id", team_id: "team-1", }; expect( - buildSavedJevConnectionTestRequest( - format === "json" ? JSON.stringify(config) : config, - "strong", - "saved-router", - "team-1", - ), + buildSavedJevConnectionTestRequest(format === "json" ? JSON.stringify(config) : config, "saved-id", "team-1"), ).toEqual(expectedRequest); }); it.each([undefined, null, "not json", "[]", {}, { classifier_type: "llm", tiers: {} }, { classifier_type: "jev" }])( "does not build a JEV probe for invalid or other classifier configurations: %j", (config) => { - expect(buildSavedJevConnectionTestRequest(config)).toBeUndefined(); + expect(buildSavedJevConnectionTestRequest(config, "saved-id")).toBeUndefined(); }, ); it("sends the prompt with the config being edited", () => { diff --git a/ui/litellm-dashboard/src/components/add_model/build_auto_router_routing_test_request.ts b/ui/litellm-dashboard/src/components/add_model/build_auto_router_routing_test_request.ts index 022bd8ad539..4679f3c50bf 100644 --- a/ui/litellm-dashboard/src/components/add_model/build_auto_router_routing_test_request.ts +++ b/ui/litellm-dashboard/src/components/add_model/build_auto_router_routing_test_request.ts @@ -6,10 +6,10 @@ export const JEV_CONNECTION_TEST_PROMPT = "What is 2 plus 2?"; export const buildSavedJevConnectionTestRequest = ( rawConfig: unknown, - defaultModel?: string, - routerName?: string, + savedModelId?: string, teamId?: string, ): AutoRouterRoutingTestRequest | undefined => { + if (!savedModelId) return undefined; const parsed: unknown = typeof rawConfig === "string" ? (() => { @@ -27,9 +27,8 @@ export const buildSavedJevConnectionTestRequest = ( if (!result.success) return undefined; return { prompt: JEV_CONNECTION_TEST_PROMPT, - complexity_router_config: result.data, - ...(defaultModel && { default_model: defaultModel }), - ...(routerName && { router_name: routerName }), + complexity_router_config: { ...result.data, jev_classifier_config: undefined }, + saved_model_id: savedModelId, ...(teamId && { team_id: teamId }), }; }; diff --git a/ui/litellm-dashboard/src/components/model_info_view.tsx b/ui/litellm-dashboard/src/components/model_info_view.tsx index 4e5ba81f2a4..7641a78cc6b 100644 --- a/ui/litellm-dashboard/src/components/model_info_view.tsx +++ b/ui/litellm-dashboard/src/components/model_info_view.tsx @@ -849,8 +849,7 @@ export default function ModelInfoView({ targets={autoRouterTestTargets} jevRequest={buildSavedJevConnectionTestRequest( (localModelData ?? modelData)?.litellm_params?.complexity_router_config, - (localModelData ?? modelData)?.litellm_params?.complexity_router_default_model, - (localModelData ?? modelData)?.model_name, + (localModelData ?? modelData)?.model_info?.id, (localModelData ?? modelData)?.model_info?.team_id, )} /> diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index 83378f984e6..b05cb48eddc 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -2327,6 +2327,7 @@ export const testModelGroupConnection = async ( export interface AutoRouterRoutingTestRequest { prompt: string; complexity_router_config: ComplexityRouterConfigPayload | Record; + saved_model_id?: string; default_model?: string; router_name?: string; team_id?: string; diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index a2a6f553da5..1f3911b3cff 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -24214,6 +24214,11 @@ export interface components { * @default auto_router_routing_test */ router_name: string; + /** + * Saved Model Id + * @description Test this saved deployment's server-side configuration instead of the supplied config and default model + */ + saved_model_id?: string | null; /** * System * @description The top-level system prompt an Anthropic /v1/messages body carries beside its messages