mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-13 23:11:40 +00:00
Merge pull request #40839 from BerriAI/litellm_compression-inheritance
fix(router): restore compression inheritance when clearing overrides
This commit is contained in:
commit
15bcc5f6eb
7 changed files with 198 additions and 12 deletions
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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 });
|
||||
|
||||
|
|
|
|||
|
|
@ -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({});
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
};
|
||||
|
|
|
|||
|
|
@ -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 () => {
|
||||
|
|
@ -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,
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue