Merge pull request #40839 from BerriAI/litellm_compression-inheritance

fix(router): restore compression inheritance when clearing overrides
This commit is contained in:
yuneng-jiang 2026-09-12 09:59:59 -07:00 committed by GitHub
commit 15bcc5f6eb
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 198 additions and 12 deletions

View file

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

View file

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

View file

@ -20,8 +20,7 @@ const NONE_OPTION: SearchSelectOption = { label: "None (no compression)", value:
const CompressionControls: React.FC<CompressionControlsProps> = ({ 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 });

View file

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

View file

@ -30,6 +30,8 @@ export interface AutoRouterCompressionLitellmParams {
auto_router_model_compression?: string;
}
type AutoRouterCompressionPatch = Partial<Record<keyof AutoRouterCompressionLitellmParams, string | null>>;
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);
};

View file

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

View file

@ -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<EditAutoRouterModalProps> = ({
...modelData.litellm_params,
complexity_router_config: updatedConfig,
complexity_router_default_model: defaultModel,
...buildAutoRouterCompressionParams(autoRouterCompression),
...buildAutoRouterCompressionPatch(autoRouterCompression, modelData.litellm_params ?? {}),
};
const updatedModelInfo = {
...modelData.model_info,