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 6cea2a946e4..5c65c2f9ba4 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,24 +7,24 @@ 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 -from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler +from litellm.proxy import proxy_server from litellm.proxy._types import ( LitellmUserRoles, ProxyErrorTypes, ProxyException, UserAPIKeyAuth, ) -from litellm.proxy import proxy_server from litellm.proxy.management_endpoints.auto_router_endpoints import ( preview_auto_router_routing, ) from litellm.router import Router -from litellm.router_strategy.complexity_router import complexity_router as complexity_module from litellm.types.management_endpoints.auto_router_endpoints import ( AutoRouterBenchmarksResponse, AutoRouterRoutingTestRequest, @@ -429,70 +429,6 @@ async def test_a_key_over_its_budget_cannot_run_a_classifier_config(monkeypatch: assert calls == [] -@pytest.mark.asyncio -@pytest.mark.parametrize("denial", ["key", "team", "budget", None]) -async def test_jev_test_routing_authorizes_paid_evaluation_before_contacting_typesafe( - monkeypatch: pytest.MonkeyPatch, denial: str | None -) -> None: - router: Final = RecordingRouter("SIMPLE") - monkeypatch.setattr(proxy_server, "llm_router", router) - monkeypatch.setenv("TYPESAFE_API_KEY", "test") - monkeypatch.setenv("TYPESAFE_API_BASE", "https://typesafe.test") - models: Final = ["cheap-model", "typesafe/jev-latest"] - actor: Final = UserAPIKeyAuth( - user_role=LitellmUserRoles.PROXY_ADMIN, - api_key="sk-jev-test", - user_id="admin", - models=["cheap-model"] if denial == "key" else models, - team_id="jev-test-team" if denial == "team" else None, - team_models=["cheap-model"] if denial == "team" else models, - max_budget=1, - spend=1 if denial == "budget" else 0, - ) - with respx.mock(assert_all_called=False) as http: - handler: Final = AsyncHTTPHandler() - handler.client = httpx.AsyncClient(transport=httpx.MockTransport(http.async_handler)) - - def http_client(_provider: object) -> AsyncHTTPHandler: - return handler - - monkeypatch.setattr(complexity_module, "get_async_httpx_client", http_client) - evaluation: Final = http.post("https://typesafe.test/v1/systemone").mock( - return_value=httpx.Response( - 200, - json={ - "answers": { - "tier": {"type": "choice", "choice": "SIMPLE", "confidence": 1, "probabilities": {"SIMPLE": 1}} - } - }, - ) - ) - call: Final = preview_auto_router_routing( - http_request=ROUTING_HTTP_REQUEST, - data=_request("small deterministic ask", classifier_type="jev", jev_classifier_config={}), - user_api_key_dict=actor, - ) - if denial is not None: - with pytest.raises(ProxyException) as exc: - await call - assert ( - exc.value.type - == { - "key": ProxyErrorTypes.key_model_access_denied, - "team": ProxyErrorTypes.team_model_access_denied, - "budget": ProxyErrorTypes.budget_exceeded, - }[denial] - ) - assert evaluation.call_count == 0 - else: - response: Final = await call - assert response.routing_decision["cause"] == "jev_classifier" - assert response.routed_model == "cheap-model" - assert evaluation.call_count == 1 - assert router.recorded_calls == [] - await handler.client.aclose() - - @pytest.mark.asyncio async def test_a_heuristic_config_does_not_need_a_budget(monkeypatch: pytest.MonkeyPatch): import litellm.proxy.proxy_server as proxy_server @@ -2352,6 +2288,70 @@ async def test_list_shadow_eval_jobs_collapses_legs_into_jobs_newest_first(monke assert group_reads == [] +@pytest.mark.asyncio +@pytest.mark.parametrize("denial", ["key", "team", "budget", None]) +async def test_jev_test_routing_authorizes_paid_evaluation_before_contacting_typesafe( + monkeypatch: pytest.MonkeyPatch, denial: str | None +) -> None: + router: Final = RecordingRouter("SIMPLE") + monkeypatch.setattr(proxy_server, "llm_router", router) + monkeypatch.setenv("TYPESAFE_API_KEY", "test") + monkeypatch.setenv("TYPESAFE_API_BASE", "https://typesafe.test") + models: Final = ["cheap-model", "typesafe/jev-latest"] + actor: Final = UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-jev-test", + user_id="admin", + models=["cheap-model"] if denial == "key" else models, + team_id="jev-test-team" if denial == "team" else None, + team_models=["cheap-model"] if denial == "team" else models, + max_budget=1, + spend=1 if denial == "budget" else 0, + ) + 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://typesafe.test/v1/systemone").mock( + return_value=httpx.Response( + 200, + json={ + "answers": { + "tier": {"type": "choice", "choice": "SIMPLE", "confidence": 1, "probabilities": {"SIMPLE": 1}} + } + }, + ) + ) + call: Final = preview_auto_router_routing( + http_request=ROUTING_HTTP_REQUEST, + data=_request("small deterministic ask", classifier_type="jev", jev_classifier_config={}), + user_api_key_dict=actor, + ) + if denial is not None: + with pytest.raises(ProxyException) as exc: + await call + assert ( + exc.value.type + == { + "key": ProxyErrorTypes.key_model_access_denied, + "team": ProxyErrorTypes.team_model_access_denied, + "budget": ProxyErrorTypes.budget_exceeded, + }[denial] + ) + assert evaluation.call_count == 0 + else: + response: Final = await call + assert response.routing_decision["cause"] == "jev_classifier" + assert response.routed_model == "cheap-model" + assert evaluation.call_count == 1 + 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 8f0ad88eb65..2a00e8bb45e 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 @@ -7,14 +7,14 @@ import { buildSavedJevConnectionTestRequest, JEV_CONNECTION_TEST_PROMPT, } from "./build_auto_router_routing_test_request"; -import { buildComplexityRouterConfig } from "./build_complexity_router_config"; +import { buildComplexityRouterConfig, type BuildComplexityRouterConfigParams } from "./build_complexity_router_config"; vi.mock( "@/app/(dashboard)/hooks/autoRouter/useComplexityScorerDefaults", async () => await import("../../../tests/mocks/complexityScorerDefaults"), ); -const config = buildComplexityRouterConfig({ +const configParams: BuildComplexityRouterConfigParams = { classifierType: "jev", jevClassifierConfig: { model: "jev-latest", timeout_ms: 3000 }, tiers: { SIMPLE: ["fast"], MEDIUM: ["mid"], COMPLEX: ["strong"], REASONING: ["reasoner"] }, @@ -43,7 +43,8 @@ const config = buildComplexityRouterConfig({ tierDistancePenalty: 0.5, adaptiveEligible: "all", returnRawModelName: false, -}); +}; +const config = buildComplexityRouterConfig(configParams); const request = buildSavedJevConnectionTestRequest(JSON.stringify(config), "fast", "my-router"); const targets = buildAutoRouterTestTargets({ tiers: Object.entries(config.tiers), @@ -92,12 +93,13 @@ describe("JEV network probes", () => { }), ); const routingCall = fetchMock.mock.calls.find(([url]) => String(url).endsWith("/auto_router/test_routing")); - expect(JSON.parse(String(routingCall?.[1]?.body))).toEqual({ + const expectedRequest = { prompt: JEV_CONNECTION_TEST_PROMPT, complexity_router_config: config, default_model: "fast", router_name: "my-router", - }); + }; + expect(JSON.parse(String(routingCall?.[1]?.body))).toEqual(expectedRequest); expect(fetchMock).toHaveBeenCalledTimes(5); expect(screen.getAllByTestId("test-status-success")).toHaveLength(4); expect(screen.getByRole("status", { name: "JEV connection" })).toHaveTextContent( 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 2aa02e40b5f..fba4ca47e00 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 @@ -29,6 +29,13 @@ describe("buildAutoRouterRoutingTestRequest", () => { fallback_tier: "DEEP", classifier_context_window_size: 4, }; + const expectedRequest = { + prompt: JEV_CONNECTION_TEST_PROMPT, + complexity_router_config: config, + default_model: "strong", + router_name: "saved-router", + team_id: "team-1", + }; expect( buildSavedJevConnectionTestRequest( format === "json" ? JSON.stringify(config) : config, @@ -36,13 +43,7 @@ describe("buildAutoRouterRoutingTestRequest", () => { "saved-router", "team-1", ), - ).toEqual({ - prompt: JEV_CONNECTION_TEST_PROMPT, - complexity_router_config: config, - default_model: "strong", - router_name: "saved-router", - team_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", diff --git a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts index e03ec22b79f..88a0cebd506 100644 --- a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts +++ b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts @@ -76,7 +76,7 @@ describe("buildComplexityRouterConfig", () => { }); it.each([false, true])("serializes JEV with shared context and no LLM config, custom tiers: %s", (custom) => { - const config = buildComplexityRouterConfig({ + const params: BuildComplexityRouterConfigParams = { ...baseParams, classifierType: "jev", jevClassifierConfig: { @@ -102,15 +102,17 @@ describe("buildComplexityRouterConfig", () => { fallback_tier_id: "quick", }, }), - }); + }; + const config = buildComplexityRouterConfig(params); expect(config.classifier_type).toBe("jev"); - expect(config.jev_classifier_config).toEqual({ + const expectedJevConfig = { model: "jev-test", timeout_ms: 4500, instructions: "Choose the configured tier", circuit_breaker_enabled: false, circuit_breaker_cooldown_seconds: 12.5, - }); + }; + expect(config.jev_classifier_config).toEqual(expectedJevConfig); expect(config.classifier_context_window_size).toBe(4); expect(config.classifier_context_budget_chars).toBe(2000); expect(config.classifier_context_include_assistant_turns).toBe(true); @@ -133,12 +135,13 @@ describe("buildComplexityRouterConfig", () => { jevClassifierConfig: { model: "jev-latest", timeout_ms: 3000, instructions: " " }, }); expect(jev.jev_classifier_config).toEqual({ model: "jev-latest", timeout_ms: 3000 }); - const llm = buildComplexityRouterConfig({ + const llmParams: BuildComplexityRouterConfigParams = { ...baseParams, classifierType: "llm", classifierLlmConfig: { model: "judge", timeout_ms: 1000 }, jevClassifierConfig: jev.jev_classifier_config, - }); + }; + const llm = buildComplexityRouterConfig(llmParams); expect(llm).not.toHaveProperty("jev_classifier_config"); }); diff --git a/ui/litellm-dashboard/src/components/add_model/classifier_type_transition.test.ts b/ui/litellm-dashboard/src/components/add_model/classifier_type_transition.test.ts index 4a0b29ecee3..a26b39c2980 100644 --- a/ui/litellm-dashboard/src/components/add_model/classifier_type_transition.test.ts +++ b/ui/litellm-dashboard/src/components/add_model/classifier_type_transition.test.ts @@ -25,7 +25,7 @@ describe("transitionClassifierType", () => { adaptive: true, }; const jev = transitionClassifierType(initial, "jev"); - expect(jev).toMatchObject({ + const expectedJevConfig = { classifier_type: "jev", jev_classifier_config: { model: "jev-latest", timeout_ms: 3000 }, classifier_context_window_size: 8, @@ -36,7 +36,8 @@ describe("transitionClassifierType", () => { enable_non_reasoning_tier: true, plan_mode_min_tier: "NON_REASONING", tiers: initial.tiers, - }); + }; + expect(jev).toMatchObject(expectedJevConfig); expect(jev.classifier_llm_config).toBeUndefined(); expect(jev.classification_prompt).toBeUndefined(); expect(jev.classification_examples).toBeUndefined(); diff --git a/ui/litellm-dashboard/src/components/add_model/jev_classifier_config.ts b/ui/litellm-dashboard/src/components/add_model/jev_classifier_config.ts index a1481c9c2e8..478c763351c 100644 --- a/ui/litellm-dashboard/src/components/add_model/jev_classifier_config.ts +++ b/ui/litellm-dashboard/src/components/add_model/jev_classifier_config.ts @@ -1,6 +1,6 @@ import { z } from "zod"; -export const jevClassifierConfigSchema = z.object({ +const jevClassifierConfigFields = { model: z.string().trim().min(1).default("jev-latest"), timeout_ms: z.number().int().positive().default(3000), instructions: z @@ -9,7 +9,9 @@ export const jevClassifierConfigSchema = z.object({ .transform((value) => value ?? undefined), circuit_breaker_enabled: z.boolean().optional(), circuit_breaker_cooldown_seconds: z.number().finite().positive().optional(), -}); +}; + +export const jevClassifierConfigSchema = z.object(jevClassifierConfigFields); export type JevClassifierConfig = z.infer; diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts b/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts index 2a3804b0307..02387dcf759 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts +++ b/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts @@ -88,14 +88,15 @@ describe("buildUpdatedComplexityRouterConfig keyword matching", () => { expect(hydrated.classifier_llm_config).toBeUndefined(); expect(hydrated.jev_classifier_config).toEqual(stored.jev_classifier_config); const saved = buildUpdatedComplexityRouterConfig(stored, hydrated); - expect(saved).toMatchObject({ + const expectedSavedConfig = { classifier_type: "jev", jev_classifier_config: stored.jev_classifier_config, classifier_context_window_size: 7, classifier_context_budget_chars: 9000, classifier_context_include_assistant_turns: true, some_future_backend_key: { nested: true }, - }); + }; + expect(saved).toMatchObject(expectedSavedConfig); expect(saved).not.toHaveProperty("classifier_llm_config"); const reloaded = hydrateComplexityRouterConfig(saved, undefined); expect(reloaded.jev_classifier_config).toEqual(hydrated.jev_classifier_config); diff --git a/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts b/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts index 442dd974368..d9e83ab850f 100644 --- a/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts +++ b/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts @@ -694,12 +694,13 @@ describe("autorouter_presets", () => { classifier_context_window_size: 6, }; const prefill = buildPresetPrefill(config, groupsOnly(["fast"])); - expect(prefill.complexityRouterConfig).toMatchObject({ + const expectedJevConfig = { classifier_type: "jev", jev_classifier_config: config.jev_classifier_config, classifier_context_window_size: 6, classifier_llm_config: undefined, - }); + }; + expect(prefill.complexityRouterConfig).toMatchObject(expectedJevConfig); const llmConfig = { ...config, classifier_type: "llm" as const }; const llmPrefill = buildPresetPrefill(llmConfig, groupsOnly(["fast"])); expect(llmPrefill.complexityRouterConfig.jev_classifier_config).toBeUndefined();