From 7e5cf9d7e7d07fd270835ab7e1408b4d1d4604e7 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Fri, 11 Sep 2026 23:24:16 -0700 Subject: [PATCH] fix(router): restore compression inheritance when clearing overrides --- .../model_management_endpoints.py | 9 ++ .../test_model_management_endpoints.py | 95 ++++++++++++++++++- .../add_model/CompressionControls.tsx | 3 +- .../buildAutoRouterCompression.test.ts | 13 +++ .../add_model/buildAutoRouterCompression.ts | 17 ++++ ...it_auto_router_modal.integration.test.tsx} | 69 ++++++++++++-- .../edit_auto_router_modal.tsx | 4 +- 7 files changed, 198 insertions(+), 12 deletions(-) rename ui/litellm-dashboard/src/components/edit_auto_router/{edit_auto_router_modal.test.tsx => edit_auto_router_modal.integration.test.tsx} (94%) diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index 94d2b773e14..8484279c69a 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -763,6 +763,15 @@ def update_db_model(db_model: Deployment, updated_patch: updateDeployment) -> Pr if field in SPECIAL_MODEL_INFO_PARAMS and getattr(updated_patch.litellm_params, field) is None: merged_litellm_params.pop(field, None) merged_model_info.pop(field, None) + elif ( + field + in ( + "auto_router_routing_compression", + "auto_router_model_compression", + ) + and getattr(updated_patch.litellm_params, field) is None + ): + merged_litellm_params.pop(field, None) if updated_patch.model_info: for field in updated_patch.model_info.model_fields_set: if field in SPECIAL_MODEL_INFO_PARAMS and getattr(updated_patch.model_info, field) is None: diff --git a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py index 5325e069813..c66095dc1fd 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py @@ -3,7 +3,7 @@ import asyncio import contextlib import json from collections.abc import Mapping -from typing import Dict, Optional +from typing import Dict, Final, Optional from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -3290,6 +3290,99 @@ def _build_db_model_with_pricing(): ) +class TestUpdateDBModelCompression: + @pytest.mark.parametrize( + "compression_patch, expected", + [ + ( + {}, + { + "auto_router_routing_compression": "routing-compressor", + "auto_router_model_compression": "model-compressor", + }, + ), + ({"auto_router_routing_compression": None}, {"auto_router_model_compression": "model-compressor"}), + ({"auto_router_model_compression": None}, {"auto_router_routing_compression": "routing-compressor"}), + ( + {"auto_router_routing_compression": "none", "auto_router_model_compression": "none"}, + {"auto_router_routing_compression": "none", "auto_router_model_compression": "none"}, + ), + ( + { + "auto_router_routing_compression": "new-compressor", + "auto_router_model_compression": "new-compressor", + }, + { + "auto_router_routing_compression": "new-compressor", + "auto_router_model_compression": "new-compressor", + }, + ), + ], + ) + def test_compression_patch_preserves_omissions_and_explicit_choices( + self, monkeypatch: pytest.MonkeyPatch, compression_patch: dict[str, str | None], expected: dict[str, str] + ): + from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper + from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model + + monkeypatch.setenv("LITELLM_SALT_KEY", "synthetic-compression-salt") + result: Final = update_db_model( + db_model=Deployment( + model_name="synthetic-router", + litellm_params=LiteLLM_Params( + model="auto_router/complexity_router", + auto_router_routing_compression=encrypt_value_helper("routing-compressor"), + auto_router_model_compression=encrypt_value_helper("model-compressor"), + ), + model_info=ModelInfo(id="compression-router"), + ), + updated_patch=updateDeployment.model_validate({"litellm_params": compression_patch}), + ) + params: Final = json.loads(result["litellm_params"]) + assert { + key: decrypt_value_helper(value=val, key=key) + for key, val in params.items() + if key in ("auto_router_routing_compression", "auto_router_model_compression") + } == expected + + def test_explicit_compression_clear_removes_both_saved_overrides(self): + from litellm.proxy.guardrails.auto_router_compression import policy_from_litellm_params + from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model + + db_model: Final = Deployment( + model_name="synthetic-router", + litellm_params=LiteLLM_Params( + model="auto_router/complexity_router", + auto_router_routing_compression="routing-compressor", + auto_router_model_compression="model-compressor", + api_base="http://127.0.0.1:9999/v1", + temperature=0, + ), + model_info=ModelInfo(id="compression-router", team_id="synthetic-team"), + ) + result: Final = update_db_model( + db_model=db_model, + updated_patch=updateDeployment.model_validate( + { + "litellm_params": { + "auto_router_routing_compression": None, + "auto_router_model_compression": None, + "api_base": None, + }, + "model_info": {"team_id": None}, + } + ), + ) + + params: Final = json.loads(result["litellm_params"]) + assert "auto_router_routing_compression" not in params + assert "auto_router_model_compression" not in params + assert policy_from_litellm_params(params) is None + assert params["api_base"] == "http://127.0.0.1:9999/v1" + assert params["temperature"] == 0 + assert json.loads(result["model_info"])["team_id"] == "synthetic-team" + + class TestUpdateDBModelClearPricing: """Sending an explicit `null` for a pricing field must remove it from both `litellm_params` and `model_info` (SPECIAL_MODEL_INFO_PARAMS are mirrored diff --git a/ui/litellm-dashboard/src/components/add_model/CompressionControls.tsx b/ui/litellm-dashboard/src/components/add_model/CompressionControls.tsx index e67583febf3..2f28e755a1d 100644 --- a/ui/litellm-dashboard/src/components/add_model/CompressionControls.tsx +++ b/ui/litellm-dashboard/src/components/add_model/CompressionControls.tsx @@ -20,8 +20,7 @@ const NONE_OPTION: SearchSelectOption = { label: "None (no compression)", value: const CompressionControls: React.FC = ({ value, onChange }) => { const { routing, sameAsRouting, model } = value; - const onRoutingChange = (newRouting: string | undefined) => - onChange({ ...value, routing: newRouting, sameAsRouting: newRouting === undefined ? true : sameAsRouting }); + const onRoutingChange = (newRouting: string | undefined) => onChange({ ...value, routing: newRouting }); const onSameAsRoutingChange = (newSameAsRouting: boolean) => onChange({ ...value, sameAsRouting: newSameAsRouting }); const onModelChange = (newModel: string | undefined) => onChange({ ...value, model: newModel }); diff --git a/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.test.ts b/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.test.ts index 20d6af50d18..a3bd882243e 100644 --- a/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.test.ts +++ b/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.test.ts @@ -1,10 +1,23 @@ import { buildAutoRouterCompressionParams, + buildAutoRouterCompressionPatch, DEFAULT_AUTO_ROUTER_COMPRESSION, hydrateAutoRouterCompression, NO_COMPRESSION, } from "./buildAutoRouterCompression"; +describe("buildAutoRouterCompressionPatch", () => { + it.each([ + {}, + { auto_router_routing_compression: "routing-compressor" }, + { auto_router_model_compression: "model-compressor" }, + { auto_router_routing_compression: "none", auto_router_model_compression: "none" }, + { auto_router_routing_compression: "routing-compressor", auto_router_model_compression: "model-compressor" }, + ])("should preserve the exact stored fields on an untouched save: %j", (stored) => { + expect(buildAutoRouterCompressionPatch(hydrateAutoRouterCompression(stored), stored)).toEqual({}); + }); +}); + describe("buildAutoRouterCompressionParams", () => { it("omits both keys when routing was never configured", () => { expect(buildAutoRouterCompressionParams(DEFAULT_AUTO_ROUTER_COMPRESSION)).toEqual({}); diff --git a/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.ts b/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.ts index 6f401a12865..c3503f78afb 100644 --- a/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.ts +++ b/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.ts @@ -30,6 +30,8 @@ export interface AutoRouterCompressionLitellmParams { auto_router_model_compression?: string; } +type AutoRouterCompressionPatch = Partial>; + export const DEFAULT_AUTO_ROUTER_COMPRESSION: AutoRouterCompressionState = { routing: undefined, sameAsRouting: true, @@ -67,3 +69,18 @@ export const hydrateAutoRouterCompression = (litellmParams: { const sameAsRouting = model === routing; return { routing, sameAsRouting, model: sameAsRouting ? undefined : model }; }; + +export const buildAutoRouterCompressionPatch = ( + state: AutoRouterCompressionState, + stored: AutoRouterCompressionPatch, +): AutoRouterCompressionPatch => { + const initial = hydrateAutoRouterCompression(stored); + const modelUnchanged = state.sameAsRouting || state.model === initial.model; + if (state.routing === initial.routing && state.sameAsRouting === initial.sameAsRouting && modelUnchanged) { + return {}; + } + if (state.routing === undefined) { + return { auto_router_routing_compression: null, auto_router_model_compression: null }; + } + return buildAutoRouterCompressionParams(state); +}; diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.tsx b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.integration.test.tsx similarity index 94% rename from ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.tsx rename to ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.integration.test.tsx index 3367c810991..277a3825d5d 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.tsx +++ b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.integration.test.tsx @@ -1064,6 +1064,55 @@ describe("EditAutoRouterModal prompt compression", () => { />, ); + it("should clear both saved compression overrides when inheritance is selected", async () => { + const user = userEvent.setup(); + renderWithStoredCompression({ + auto_router_routing_compression: "routing-compressor", + auto_router_model_compression: "model-compressor", + }); + + await user.click(await screen.findByText("Advanced: Compression")); + await user.click(screen.getAllByRole("button", { name: "Clear", exact: true })[0]); + await user.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => + expect(modelPatchUpdateCall).toHaveBeenCalledWith( + "token", + expect.objectContaining({ + litellm_params: expect.objectContaining({ + model: "auto_router/complexity_router", + auto_router_routing_compression: null, + auto_router_model_compression: null, + }), + }), + "auto-1", + ), + ); + }); + + it("should discard a cancelled clear and preserve compression when the saved choice is restored", async () => { + const user = userEvent.setup(); + const stored = { auto_router_routing_compression: "none", auto_router_model_compression: "model-compressor" }; + const view = renderWithStoredCompression(stored); + + await user.click(await screen.findByText("Advanced: Compression")); + await user.click(screen.getAllByRole("button", { name: "Clear", exact: true })[0]); + await user.click(screen.getByRole("button", { name: "Cancel", exact: true })); + expect(modelPatchUpdateCall).not.toHaveBeenCalled(); + view.unmount(); + + renderWithStoredCompression(stored); + await user.click(await screen.findByText("Advanced: Compression")); + expect(screen.getByRole("combobox", { name: "Routing decision compression" })).toHaveValue("None (no compression)"); + await user.click(screen.getAllByRole("button", { name: "Clear", exact: true })[0]); + await user.click(screen.getByRole("combobox", { name: "Routing decision compression" })); + await user.click(screen.getByRole("option", { name: "None (no compression)" })); + await user.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => expect(modelPatchUpdateCall).toHaveBeenCalled()); + expect(savedLitellmParams()).toMatchObject(stored); + }); + it("leaves both compression keys out of an untouched save when none were stored", async () => { const user = userEvent.setup(); renderWithStoredCompression(); @@ -1075,18 +1124,24 @@ describe("EditAutoRouterModal prompt compression", () => { expect(savedLitellmParams()).not.toHaveProperty("auto_router_model_compression"); }); - it("preserves a stored same-as-routing compression through an untouched open-and-save", async () => { + it.each([ + { auto_router_routing_compression: "headroom-a", auto_router_model_compression: "headroom-a" }, + { auto_router_routing_compression: "routing-compressor" }, + { auto_router_model_compression: "model-compressor" }, + ])("should preserve the exact stored compression fields through an untouched save: %j", async (stored) => { const user = userEvent.setup(); - renderWithStoredCompression({ - auto_router_routing_compression: "headroom-a", - auto_router_model_compression: "headroom-a", - }); + renderWithStoredCompression(stored); await user.click(await screen.findByRole("button", { name: /save changes/i })); await waitFor(() => expect(modelPatchUpdateCall).toHaveBeenCalled()); - expect(savedLitellmParams()?.auto_router_routing_compression).toBe("headroom-a"); - expect(savedLitellmParams()?.auto_router_model_compression).toBe("headroom-a"); + expect( + Object.fromEntries( + Object.entries(savedLitellmParams()).filter( + ([key]) => key === "auto_router_routing_compression" || key === "auto_router_model_compression", + ), + ), + ).toEqual(stored); }); it("shows a stored different-compression choice as Use a different compression, not Same", async () => { diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx index 86f18ee9b12..5ae9b005f20 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx +++ b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx @@ -44,7 +44,7 @@ import { KeywordTierRule } from "../add_model/KeywordTierRules"; import { DEFAULT_MATCH_THRESHOLD } from "../add_model/SemanticKeywordMatching"; import { type AutoRouterCompressionState, - buildAutoRouterCompressionParams, + buildAutoRouterCompressionPatch, DEFAULT_AUTO_ROUTER_COMPRESSION, hydrateAutoRouterCompression, } from "../add_model/buildAutoRouterCompression"; @@ -679,7 +679,7 @@ const EditAutoRouterModal: React.FC = ({ ...modelData.litellm_params, complexity_router_config: updatedConfig, complexity_router_default_model: defaultModel, - ...buildAutoRouterCompressionParams(autoRouterCompression), + ...buildAutoRouterCompressionPatch(autoRouterCompression, modelData.litellm_params ?? {}), }; const updatedModelInfo = { ...modelData.model_info,