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

fix(auto-router): resolve saved JEV probes on the server

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Moe Khalil 2026-09-20 00:54:17 +00:00
parent 401baf32c3
commit 97c54e278e
9 changed files with 187 additions and 40 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

68
litellm/proxy/management_endpoints/auto_router_endpoints.py
View file

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

5
litellm/types/management_endpoints/auto_router_endpoints.py
View file

@ -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",

93
tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py
View file

@ -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

13
ui/litellm-dashboard/src/components/add_model/JevConnectionTest.integration.test.tsx
View file

@ -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);

30
ui/litellm-dashboard/src/components/add_model/build_auto_router_routing_test_request.test.ts
View file

@ -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", () => {

9
ui/litellm-dashboard/src/components/add_model/build_auto_router_routing_test_request.ts
View file

@ -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 }),
};
};

3
ui/litellm-dashboard/src/components/model_info_view.tsx
View file

@ -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,
)}
/>

1
ui/litellm-dashboard/src/components/networking.tsx
View file

@ -2327,6 +2327,7 @@ export const testModelGroupConnection = async (
export interface AutoRouterRoutingTestRequest {
prompt: string;
complexity_router_config: ComplexityRouterConfigPayload | Record<string, unknown>;
saved_model_id?: string;
default_model?: string;
router_name?: string;
team_id?: string;

5
ui/litellm-dashboard/src/lib/http/schema.d.ts generated vendored
View file

@ -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