From dd60b7e40f44d7687ad4577332d6f57fc21d7af3 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Fri, 4 Sep 2026 16:16:38 -0700 Subject: [PATCH] feat(auto-router): decouple compression between the routing decision and the model call An auto router marker deployment can now set auto_router_routing_compression and auto_router_model_compression in its litellm_params, naming the compression guardrail each hop should use (or "none" for no compression on that hop). Neither key set means the request's own compression guardrails keep applying to both hops unchanged. Backend: Router.async_pre_routing_hook resolves the marker's policy and compresses a copy of the messages for the routing decision only when the policy differs from what the model call already got; when both hops share the same compression, it reuses what the ordinary pre-call guardrail pipeline already produced instead of compressing twice. The proxy layer suppresses every other compression guardrail once a policy is engaged and arms the model-side guardrail even when it is not default_on. UI: the auto router's Detailed Configuration gains an Advanced: Compression section with a routing-decision selector and a same/different toggle for the model call, matching the same/different address pattern. --- litellm/constants.py | 4 + litellm/integrations/custom_guardrail.py | 18 ++ litellm/proxy/common_request_processing.py | 7 + .../guardrails/auto_router_compression.py | 211 +++++++++++++ litellm/router.py | 53 +++- litellm/types/router.py | 4 + litellm/types/utils.py | 2 + .../integrations/test_custom_guardrail.py | 44 +++ .../test_auto_router_compression.py | 285 ++++++++++++++++++ .../proxy/test_common_request_processing.py | 54 ++++ tests/test_litellm/test_router.py | 151 ++++++++++ .../add_model/ComplexityRouterConfig.tsx | 34 +++ .../add_model/CompressionControls.tsx | 93 ++++++ .../add_model/add_auto_router_tab.test.tsx | 75 ++++- .../add_model/add_auto_router_tab.tsx | 11 + .../buildAutoRouterCompression.test.ts | 93 ++++++ .../add_model/buildAutoRouterCompression.ts | 52 ++++ .../handle_add_auto_router_submit.tsx | 5 +- .../edit_auto_router_modal.test.tsx | 81 +++++ .../edit_auto_router_modal.tsx | 18 ++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 8 + 21 files changed, 1297 insertions(+), 6 deletions(-) create mode 100644 litellm/proxy/guardrails/auto_router_compression.py create mode 100644 tests/test_litellm/proxy/guardrails/test_auto_router_compression.py create mode 100644 ui/litellm-dashboard/src/components/add_model/CompressionControls.tsx create mode 100644 ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.test.ts create mode 100644 ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.ts diff --git a/litellm/constants.py b/litellm/constants.py index da731cb5eb2..25fdaec20de 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -215,6 +215,10 @@ MAX_CALLBACKS: Final = get_env_int("LITELLM_MAX_CALLBACKS", 100) # so the deployment-level hook does not re-run them for the same request PRE_CALL_EXECUTED_GUARDRAILS_KEY: Final = "_pre_call_executed_guardrails" +# Metadata key listing compression guardrails an auto router's own compression +# policy suppresses for this request. See litellm.proxy.guardrails.auto_router_compression. +AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY: Final = "_auto_router_suppressed_compression_guardrails" + # Generic fallback for unknown models DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET: Final = int( os.getenv("DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET", 128) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 9bb613654e4..7f8effa2317 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -45,6 +45,7 @@ dc: Final = DualCache() from litellm.constants import ( + AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY, GUARDRAIL_SCANNED_MESSAGES_CACHE_TTL_SECONDS, PRE_CALL_EXECUTED_GUARDRAILS_KEY, ) @@ -940,6 +941,20 @@ class CustomGuardrail(CustomLogger): """ return False + def _suppressed_by_auto_router_compression(self, data: dict) -> bool: + """True when an auto router's own compression policy suppresses this guardrail. + + Set only by litellm.proxy.guardrails.auto_router_compression.arm_pre_call, never + by the caller, so a request cannot suppress its own guardrails this way. + """ + for meta_key in ("metadata", "litellm_metadata"): + meta = data.get(meta_key) + if isinstance(meta, dict): + suppressed = meta.get(AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY) + if isinstance(suppressed, list) and self.guardrail_name in suppressed: + return True + return False + def should_run_guardrail( self, data, @@ -948,6 +963,9 @@ class CustomGuardrail(CustomLogger): """ Returns True if the guardrail should be run on the event_type """ + if self._suppressed_by_auto_router_compression(data): + return False + requested_guardrails: Final = self.get_guardrail_from_metadata(data) disable_global_guardrail: Final = self.get_disable_global_guardrail(data) opted_out_global_guardrails: Final = self.get_opted_out_global_guardrails_from_metadata(data) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index f25fa46197e..534b2db3e61 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -61,6 +61,7 @@ from litellm.proxy.common_utils.sse_keepalive import ( wrap_sse_stream_with_keepalive_pings, ) from litellm.proxy.dd_span_tagger import DDSpanTagger +from litellm.proxy.guardrails.auto_router_compression import arm_pre_call as _arm_auto_router_compression from litellm.proxy.route_llm_request import route_request from litellm.proxy.utils import ProxyLogging, _check_and_merge_model_level_guardrails from litellm.router import Router @@ -2004,6 +2005,12 @@ class ProxyBaseLLMRequestProcessing: trust_client_model_info=False, ) + # An auto router with its own compression policy is authoritative for this + # request: suppress every other compression guardrail and arm whichever one + # the policy names for the model call, before those guardrails get a chance + # to run below. + self.data = await _arm_auto_router_compression(data=self.data, llm_router=llm_router) + self.data = await proxy_logging_obj.pre_call_hook( user_api_key_dict=user_api_key_dict, data=self.data, diff --git a/litellm/proxy/guardrails/auto_router_compression.py b/litellm/proxy/guardrails/auto_router_compression.py new file mode 100644 index 00000000000..c3fba937d22 --- /dev/null +++ b/litellm/proxy/guardrails/auto_router_compression.py @@ -0,0 +1,211 @@ +""" +Decouples prompt compression between an auto router's routing decision and the +model it routes to. An auto router marker deployment may set +``auto_router_routing_compression`` and/or ``auto_router_model_compression`` in its +``litellm_params`` to name the compression guardrail that hop should use, or +``"none"`` to run no compression on that hop. Neither key set means the request's +own compression guardrails (key/team/model-level, or an "Always on" guardrail) +apply to both hops unchanged, exactly as before this feature existed. + +Once either key is set, this auto router is authoritative: every other compression +guardrail is suppressed for that request, and only these two settings decide what +each hop sees. +""" + +import copy +from collections.abc import Mapping +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Final + +from litellm._logging import verbose_proxy_logger +from litellm.constants import AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY +from litellm.litellm_core_utils.core_helpers import ( + get_metadata_variable_name_from_kwargs, + get_or_create_metadata_bucket, +) +from litellm.router_utils.auto_router_model_naming import AUTO_ROUTER_MODEL_PREFIX +from litellm.types.utils import GenericGuardrailAPIInputs + +if TYPE_CHECKING: + from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.router import Router +else: + CustomGuardrail = Any + Router = Any + +COMPRESSION_GUARDRAIL_PROVIDERS: Final = frozenset({"headroom", "compresr"}) +_NO_COMPRESSION: Final = "none" + +# Metadata key stashing the pre-compression messages so a routing decision that +# names a different compression than the model call still compresses the +# original text, not whatever the model-side guardrail already rewrote it to. +AUTO_ROUTER_ROUTING_MESSAGES_SNAPSHOT_KEY: Final = "_auto_router_routing_messages_snapshot" + + +@dataclass(frozen=True, slots=True) +class AutoRouterCompressionPolicy: + """An auto router's compression choice for each hop. ``None`` means no compression.""" + + routing: str | None + model: str | None + + @property + def is_same(self) -> bool: + return self.routing == self.model + + +def _normalized_compression_choice(raw: object) -> str | None: + if not isinstance(raw, str) or not raw: + return None + return None if raw.strip().lower() == _NO_COMPRESSION else raw + + +def policy_from_litellm_params(litellm_params: Mapping[str, object]) -> AutoRouterCompressionPolicy | None: + raw_routing: Final = litellm_params.get("auto_router_routing_compression") + raw_model: Final = litellm_params.get("auto_router_model_compression") + if raw_routing is None and raw_model is None: + return None + return AutoRouterCompressionPolicy( + routing=_normalized_compression_choice(raw_routing), + model=_normalized_compression_choice(raw_model), + ) + + +def policy_for_model( + llm_router: "Router | None", model_alias: str, team_id: str | None +) -> AutoRouterCompressionPolicy | None: + """The compression policy declared by the auto router marker deployment `model_alias` resolves to. + + Mirrors the alias lookup in ``_check_and_merge_model_level_guardrails``: this runs + before routing has picked a strategy, so it takes the first marker deployment for + the alias rather than disambiguating by request tags. + """ + if llm_router is None: + return None + deployments: Final = llm_router.get_model_list(model_name=model_alias, team_id=team_id) or [] + for deployment in deployments: + litellm_params: Final = deployment.get("litellm_params") or {} + model_field = litellm_params.get("model") + if not isinstance(model_field, str) or not model_field.startswith(AUTO_ROUTER_MODEL_PREFIX): + continue + policy = policy_from_litellm_params(litellm_params) + if policy is not None: + return policy + return None + + +def _active_compression_guardrail_names() -> frozenset[str]: + """Names of every currently-active guardrail whose type is a compression guardrail.""" + import litellm + from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.proxy.guardrails.guardrail_registry import guardrail_class_registry + + compression_classes: Final = tuple( + cls for name, cls in guardrail_class_registry.items() if name in COMPRESSION_GUARDRAIL_PROVIDERS + ) + if not compression_classes: + return frozenset() + active: Final = litellm.logging_callback_manager.get_custom_loggers_for_type(callback_type=CustomGuardrail) + return frozenset( + cb.guardrail_name for cb in active if isinstance(cb, compression_classes) and cb.guardrail_name + ) + + +async def arm_pre_call(data: dict, llm_router: "Router | None") -> dict: + """Apply an auto router's compression policy, if any, before guardrails run. + + Suppresses every other compression guardrail, re-enables the model-side + guardrail the policy names (if any) even when it isn't ``default_on``, and + snapshots the pre-compression messages so the routing decision can compress + them independently of whatever the model-side guardrail does to `data`. + """ + if llm_router is None: + return data + + model_alias: Final = data.get("model") + if not isinstance(model_alias, str) or not model_alias: + return data + + # Read-only until a policy is confirmed: creating the metadata bucket for every + # request, including the vast majority with no auto-router compression policy, + # would be an unwanted side effect of merely checking for one. + metadata_key: Final = get_metadata_variable_name_from_kwargs(data) + existing_bucket: Final = data.get(metadata_key) + other_bucket: Final = data.get("metadata" if metadata_key == "litellm_metadata" else "litellm_metadata") + team_id: Final = (existing_bucket.get("user_api_key_team_id") if isinstance(existing_bucket, dict) else None) or ( + other_bucket.get("user_api_key_team_id") if isinstance(other_bucket, dict) else None + ) + + policy: Final = policy_for_model(llm_router=llm_router, model_alias=model_alias, team_id=team_id) + if policy is None: + return data + + _, metadata = get_or_create_metadata_bucket(data) + suppressed: Final = _active_compression_guardrail_names() - ({policy.model} if policy.model else set()) + if suppressed: + metadata[AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY] = sorted(suppressed) + + if policy.model is not None: + requested = metadata.get("guardrails") + if isinstance(requested, list): + if policy.model not in requested: + requested.append(policy.model) + else: + metadata["guardrails"] = [policy.model] + + from litellm.litellm_core_utils.prompt_templates.factory import resolve_structured_messages + + snapshot: Final = resolve_structured_messages(messages=data.get("messages"), request_kwargs=data) + if snapshot is not None: + metadata[AUTO_ROUTER_ROUTING_MESSAGES_SNAPSHOT_KEY] = copy.deepcopy(snapshot) + + return data + + +async def messages_for_routing( + policy: AutoRouterCompressionPolicy | None, + messages: list[dict[str, Any]] | None, + request_kwargs: Mapping[str, object], +) -> list[dict[str, Any]] | None: + """Messages to use for a routing decision, compressed per `policy.routing`. + + Returns None when there is no policy or the policy's routing side names no + compression, meaning the caller should route on whatever messages it already + has. The model call is untouched by this function either way: model-side + compression, if any, already ran as an ordinary pre-call guardrail before the + router was ever reached. + """ + if policy is None or policy.routing is None: + return None + + from litellm.proxy.common_utils.registry_read_through import ( + get_initialized_guardrail_with_read_through, + ) + + metadata_key: Final = get_metadata_variable_name_from_kwargs(request_kwargs) + metadata: Final = request_kwargs.get(metadata_key) + snapshot: Final = metadata.get(AUTO_ROUTER_ROUTING_MESSAGES_SNAPSHOT_KEY) if isinstance(metadata, dict) else None + original: Final = snapshot if isinstance(snapshot, list) else messages + if not original: + return None + + guardrail: Final = await get_initialized_guardrail_with_read_through(policy.routing) + if guardrail is None: + verbose_proxy_logger.warning( + "AutoRouter compression: guardrail '%s' not found; routing on uncompressed messages", policy.routing + ) + return None + + inputs: GenericGuardrailAPIInputs = {"structured_messages": [dict(m) for m in original]} + # A throwaway request_data: apply_guardrail writes its stats onto this dict, not + # the real request's metadata, so routing-side compression never double-counts + # against extract_compression_saved_tokens's model-savings accounting. + throwaway_request_data: Final[dict[str, object]] = { + "messages": original, + "model": request_kwargs.get("model"), + } + result: Final = await guardrail.apply_guardrail( + inputs=inputs, request_data=throwaway_request_data, input_type="request" + ) + compressed = result.get("structured_messages") + return compressed if isinstance(compressed, list) else original diff --git a/litellm/router.py b/litellm/router.py index 6c7611c6236..8149ec60ddc 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -13037,13 +13037,46 @@ class Router: ) return None + from litellm.proxy.guardrails.auto_router_compression import ( + messages_for_routing, + policy_from_litellm_params, + ) + + marker_params: Final = self._alias_marker_litellm_params(registered_model_name, selected_strategy.tags) + compression_policy: Final = policy_from_litellm_params(marker_params) if marker_params else None + # When both hops share the same compression, the model-side guardrail already + # ran in the proxy's ordinary pre-call hook and compressed `messages` in place + # (arm_pre_call armed it whether or not it is `default_on`); reuse that result + # for routing too instead of paying for a second compression call against the + # same content. + needs_independent_routing_compression: Final = compression_policy is not None and not ( + compression_policy.is_same and compression_policy.model is not None + ) + routing_messages: Final = ( + await messages_for_routing(policy=compression_policy, messages=messages, request_kwargs=request_kwargs) + if needs_independent_routing_compression + else None + ) + pre_routing_hook_response: Final = await selected_strategy.strategy.async_pre_routing_hook( model=registered_model_name, request_kwargs=request_kwargs, - messages=messages, + messages=routing_messages if routing_messages is not None else messages, input=input, specific_deployment=specific_deployment, ) + # The strategy only echoes back whatever `messages` it was handed, so a + # routing-only compression must not leak into the response: the model call + # and downstream deployment-context filtering both key off this field. + # Compared by value, not identity: PreRoutingHookResponse is a pydantic model, + # and pydantic reconstructs a validated list field rather than keeping the + # exact object passed in, even when nothing about it changed. + if ( + pre_routing_hook_response is not None + and routing_messages is not None + and pre_routing_hook_response.messages == routing_messages + ): + pre_routing_hook_response = pre_routing_hook_response.model_copy(update={"messages": messages}) self._record_routing_decision( request_kwargs=request_kwargs, routing_decision=(pre_routing_hook_response.routing_decision if pre_routing_hook_response else None), @@ -13100,9 +13133,16 @@ class Router: return pre_routing_hook_response - def _forwardable_alias_marker_params( + def _alias_marker_litellm_params( self, model: str, strategy_tags: tuple[str, ...] - ) -> tuple[tuple[str, object], ...]: + ) -> Mapping[str, object] | None: + """The auto-router marker deployment's own `litellm_params` for `model`, tag-scoped. + + Shared by `_forwardable_alias_marker_params` (forwarding api_base/api_key/... + gaps onto the routed deployment) and the auto-router compression policy lookup + (reading `auto_router_routing_compression`/`auto_router_model_compression`), so + both read the same marker row when an alias has more than one, tag-scoped marker. + """ marker_params: Final = tuple( litellm_params for idx in self.model_name_to_deployment_indices.get(model, ()) @@ -13112,7 +13152,12 @@ class Router: tag_matched: Final = tuple( params for params in marker_params if tuple(params.get("tags") or ()) == strategy_tags ) - selected: Final = tag_matched[0] if tag_matched else (marker_params[0] if marker_params else None) + return tag_matched[0] if tag_matched else (marker_params[0] if marker_params else None) + + def _forwardable_alias_marker_params( + self, model: str, strategy_tags: tuple[str, ...] + ) -> tuple[tuple[str, object], ...]: + selected: Final = self._alias_marker_litellm_params(model, strategy_tags) if selected is None: return () return tuple( diff --git a/litellm/types/router.py b/litellm/types/router.py index 7ebd50f1328..f5295d6569c 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -359,6 +359,10 @@ class GenericLiteLLMParams(CredentialLiteLLMParams, CustomPricingLiteLLMParams): auto_router_default_model: str | None = None auto_router_embedding_model: str | None = None auto_router_max_input_chars: int | None = None + # Compression policy for the two hops of a routed request. Both unset means the + # request's own compression guardrails apply to both, as they always have. + auto_router_routing_compression: str | None = None + auto_router_model_compression: str | None = None # complexity-router params complexity_router_config: dict | None = None diff --git a/litellm/types/utils.py b/litellm/types/utils.py index c1d90694f14..a1fda3d0524 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3713,6 +3713,8 @@ all_litellm_params = ( "auto_router_default_model", "auto_router_embedding_model", "auto_router_max_input_chars", + "auto_router_routing_compression", + "auto_router_model_compression", "complexity_router_config", "complexity_router_default_model", "adaptive_router_config", diff --git a/tests/test_litellm/integrations/test_custom_guardrail.py b/tests/test_litellm/integrations/test_custom_guardrail.py index 6edc2c9bf77..1fb4299cb56 100644 --- a/tests/test_litellm/integrations/test_custom_guardrail.py +++ b/tests/test_litellm/integrations/test_custom_guardrail.py @@ -518,6 +518,50 @@ class TestCustomGuardrailShouldRunGuardrail: is True ) + def test_should_run_guardrail_suppressed_by_auto_router_compression(self): + """An auto router's own compression policy can suppress an otherwise-eligible + guardrail, even one that is default_on and explicitly requested.""" + from litellm.constants import AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY + from litellm.types.guardrails import GuardrailEventHooks + + always_on = CustomGuardrail( + guardrail_name="headroom-default", + default_on=True, + event_hook=GuardrailEventHooks.pre_call, + ) + data = { + "model": "smart-router", + "metadata": { + AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY: ["headroom-default"], + }, + } + + assert ( + always_on.should_run_guardrail(data=data, event_type=GuardrailEventHooks.pre_call) + is False + ) + + def test_should_run_guardrail_suppression_list_does_not_affect_other_names(self): + from litellm.constants import AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY + from litellm.types.guardrails import GuardrailEventHooks + + always_on = CustomGuardrail( + guardrail_name="headroom-default", + default_on=True, + event_hook=GuardrailEventHooks.pre_call, + ) + data = { + "model": "smart-router", + "metadata": { + AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY: ["some-other-guardrail"], + }, + } + + assert ( + always_on.should_run_guardrail(data=data, event_type=GuardrailEventHooks.pre_call) + is True + ) + class TestApplyGuardrailCheck: def test_apply_guardrail_check_only_on_direct_implementation(self): diff --git a/tests/test_litellm/proxy/guardrails/test_auto_router_compression.py b/tests/test_litellm/proxy/guardrails/test_auto_router_compression.py new file mode 100644 index 00000000000..b2e83e75768 --- /dev/null +++ b/tests/test_litellm/proxy/guardrails/test_auto_router_compression.py @@ -0,0 +1,285 @@ +""" +Unit tests for litellm.proxy.guardrails.auto_router_compression. + +Covers: +- policy_from_litellm_params: absent keys mean no policy; the "none" sentinel + normalizes to explicit no-compression within an active policy; is_same +- policy_for_model: finds the auto-router marker deployment for an alias +- arm_pre_call: no-op without a policy; suppresses active compression guardrails; + arms the model-side guardrail even when it isn't default_on; snapshots messages +- messages_for_routing: no-op without a policy or an unset routing side; compresses + via the named guardrail's apply_guardrail; never writes stats onto the caller's + own request_kwargs (regression for double-counted compression savings) +""" + +from typing import Any + +import pytest + +from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.proxy.guardrails.auto_router_compression import ( + AUTO_ROUTER_ROUTING_MESSAGES_SNAPSHOT_KEY, + AutoRouterCompressionPolicy, + arm_pre_call, + messages_for_routing, + policy_for_model, + policy_from_litellm_params, +) +from litellm.constants import AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY +from litellm.types.utils import GenericGuardrailAPIInputs + + +class TestPolicyFromLitellmParams: + def test_neither_key_set_is_no_policy(self): + assert policy_from_litellm_params({}) is None + + def test_routing_only(self): + policy = policy_from_litellm_params({"auto_router_routing_compression": "headroom-a"}) + assert policy == AutoRouterCompressionPolicy(routing="headroom-a", model=None) + + def test_none_sentinel_normalizes_to_no_compression(self): + policy = policy_from_litellm_params( + {"auto_router_routing_compression": "headroom-a", "auto_router_model_compression": "none"} + ) + assert policy == AutoRouterCompressionPolicy(routing="headroom-a", model=None) + + def test_none_sentinel_is_case_insensitive(self): + policy = policy_from_litellm_params({"auto_router_routing_compression": "NONE"}) + assert policy == AutoRouterCompressionPolicy(routing=None, model=None) + + def test_is_same_true_for_matching_names(self): + policy = policy_from_litellm_params( + {"auto_router_routing_compression": "x", "auto_router_model_compression": "x"} + ) + assert policy.is_same is True + + def test_is_same_false_for_different_names(self): + policy = policy_from_litellm_params( + {"auto_router_routing_compression": "x", "auto_router_model_compression": "y"} + ) + assert policy.is_same is False + + def test_is_same_true_when_both_no_compression(self): + policy = policy_from_litellm_params( + {"auto_router_routing_compression": "none", "auto_router_model_compression": "none"} + ) + assert policy.is_same is True + + +class _FakeRouter: + """Minimal stand-in for litellm.Router.get_model_list, for policy_for_model.""" + + def __init__(self, deployments: list[dict[str, Any]]): + self._deployments = deployments + + def get_model_list(self, model_name, team_id=None): + return [d for d in self._deployments if d.get("model_name") == model_name] + + +class TestPolicyForModel: + def test_no_router_returns_none(self): + assert policy_for_model(llm_router=None, model_alias="smart-router", team_id=None) is None + + def test_no_marker_deployment_returns_none(self): + router = _FakeRouter( + [{"model_name": "smart-router", "litellm_params": {"model": "openai/gpt-4o-mini"}}] + ) + assert policy_for_model(llm_router=router, model_alias="smart-router", team_id=None) is None + + def test_marker_deployment_without_policy_returns_none(self): + router = _FakeRouter( + [{"model_name": "smart-router", "litellm_params": {"model": "auto_router/complexity_router"}}] + ) + assert policy_for_model(llm_router=router, model_alias="smart-router", team_id=None) is None + + def test_marker_deployment_with_policy_is_found(self): + router = _FakeRouter( + [ + { + "model_name": "smart-router", + "litellm_params": { + "model": "auto_router/complexity_router", + "auto_router_routing_compression": "headroom-a", + "auto_router_model_compression": "none", + }, + } + ] + ) + policy = policy_for_model(llm_router=router, model_alias="smart-router", team_id=None) + assert policy == AutoRouterCompressionPolicy(routing="headroom-a", model=None) + + +class _RecordingCompressionGuardrail(CustomGuardrail): + """A guardrail whose apply_guardrail marks every text message as compressed.""" + + def __init__(self, guardrail_name: str): + super().__init__(guardrail_name=guardrail_name) + self.request_data_seen: list[dict] = [] + + async def apply_guardrail( + self, inputs: GenericGuardrailAPIInputs, request_data: dict, input_type: str, logging_obj=None + ) -> GenericGuardrailAPIInputs: + self.request_data_seen.append(request_data) + structured_messages = inputs.get("structured_messages") or [] + compressed = [ + {**m, "content": f"[COMPRESSED] {m.get('content')}"} for m in structured_messages + ] + return {**inputs, "structured_messages": compressed} + + +@pytest.fixture +def registered_guardrail(): + import litellm + + guardrail = _RecordingCompressionGuardrail(guardrail_name="fake-compress") + litellm.logging_callback_manager.add_litellm_callback(guardrail) + yield guardrail + litellm.logging_callback_manager.remove_callback_from_all_lists(guardrail) + + +class TestArmPreCall: + @pytest.mark.asyncio + async def test_no_router_is_noop(self): + data = {"model": "smart-router", "messages": [{"role": "user", "content": "hi"}]} + result = await arm_pre_call(data=data, llm_router=None) + assert result == data + assert "metadata" not in result + + @pytest.mark.asyncio + async def test_no_policy_does_not_create_metadata_bucket(self): + router = _FakeRouter( + [{"model_name": "smart-router", "litellm_params": {"model": "openai/gpt-4o-mini"}}] + ) + data = {"model": "smart-router", "messages": [{"role": "user", "content": "hi"}]} + result = await arm_pre_call(data=data, llm_router=router) + assert "metadata" not in result + assert "litellm_metadata" not in result + + @pytest.mark.asyncio + async def test_policy_suppresses_active_compression_guardrails(self, monkeypatch): + from litellm.proxy.guardrails import guardrail_registry + + monkeypatch.setitem( + guardrail_registry.guardrail_class_registry, "fake-provider", _RecordingCompressionGuardrail + ) + monkeypatch.setattr( + "litellm.proxy.guardrails.auto_router_compression.COMPRESSION_GUARDRAIL_PROVIDERS", + frozenset({"fake-provider"}), + ) + import litellm + + always_on = _RecordingCompressionGuardrail(guardrail_name="always-on-compression") + litellm.logging_callback_manager.add_litellm_callback(always_on) + try: + router = _FakeRouter( + [ + { + "model_name": "smart-router", + "litellm_params": { + "model": "auto_router/complexity_router", + "auto_router_routing_compression": "headroom-a", + "auto_router_model_compression": "none", + }, + } + ] + ) + data = {"model": "smart-router", "messages": [{"role": "user", "content": "hi"}]} + result = await arm_pre_call(data=data, llm_router=router) + suppressed = result["metadata"][AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY] + assert "always-on-compression" in suppressed + finally: + litellm.logging_callback_manager.remove_callback_from_all_lists(always_on) + + @pytest.mark.asyncio + async def test_model_side_guardrail_is_requested_even_when_not_default_on(self): + router = _FakeRouter( + [ + { + "model_name": "smart-router", + "litellm_params": { + "model": "auto_router/complexity_router", + "auto_router_routing_compression": "none", + "auto_router_model_compression": "headroom-b", + }, + } + ] + ) + data = {"model": "smart-router", "messages": [{"role": "user", "content": "hi"}]} + result = await arm_pre_call(data=data, llm_router=router) + assert result["metadata"]["guardrails"] == ["headroom-b"] + + @pytest.mark.asyncio + async def test_snapshots_original_messages(self): + router = _FakeRouter( + [ + { + "model_name": "smart-router", + "litellm_params": { + "model": "auto_router/complexity_router", + "auto_router_routing_compression": "headroom-a", + "auto_router_model_compression": "none", + }, + } + ] + ) + original_messages = [{"role": "user", "content": "hi"}] + data = {"model": "smart-router", "messages": original_messages} + result = await arm_pre_call(data=data, llm_router=router) + snapshot = result["metadata"][AUTO_ROUTER_ROUTING_MESSAGES_SNAPSHOT_KEY] + assert snapshot == original_messages + assert snapshot is not original_messages # a copy, not the live reference + + +class TestMessagesForRouting: + @pytest.mark.asyncio + async def test_no_policy_returns_none(self): + assert await messages_for_routing(policy=None, messages=[], request_kwargs={}) is None + + @pytest.mark.asyncio + async def test_routing_side_unset_returns_none(self): + policy = AutoRouterCompressionPolicy(routing=None, model="headroom-a") + assert await messages_for_routing(policy=policy, messages=[], request_kwargs={}) is None + + @pytest.mark.asyncio + async def test_unknown_guardrail_name_returns_none(self): + policy = AutoRouterCompressionPolicy(routing="does-not-exist", model=None) + messages = [{"role": "user", "content": "hi"}] + result = await messages_for_routing(policy=policy, messages=messages, request_kwargs={}) + assert result is None + + @pytest.mark.asyncio + async def test_compresses_via_the_named_guardrail(self, registered_guardrail): + policy = AutoRouterCompressionPolicy(routing="fake-compress", model=None) + messages = [{"role": "user", "content": "hello world"}] + result = await messages_for_routing(policy=policy, messages=messages, request_kwargs={}) + assert result == [{"role": "user", "content": "[COMPRESSED] hello world"}] + + @pytest.mark.asyncio + async def test_uses_the_snapshot_when_present(self, registered_guardrail): + policy = AutoRouterCompressionPolicy(routing="fake-compress", model=None) + snapshot = [{"role": "user", "content": "original"}] + request_kwargs = {"metadata": {AUTO_ROUTER_ROUTING_MESSAGES_SNAPSHOT_KEY: snapshot}} + # `messages` here stands in for whatever a model-side guardrail already + # rewrote `data["messages"]` to -- routing must ignore it and compress the + # pristine snapshot instead. + already_rewritten = [{"role": "user", "content": "rewritten by another guardrail"}] + result = await messages_for_routing( + policy=policy, messages=already_rewritten, request_kwargs=request_kwargs + ) + assert result == [{"role": "user", "content": "[COMPRESSED] original"}] + + @pytest.mark.asyncio + async def test_guardrail_receives_a_throwaway_request_data_not_the_real_request_kwargs( + self, registered_guardrail + ): + """Regression: a real compression guardrail writes its stats onto whatever + `request_data` dict it's given (`add_standard_logging_guardrail_information_to_ + request_data`). If that were the caller's own `request_kwargs`, routing-side + compression would double-count into extract_compression_saved_tokens, which + sums every guardrail_information entry on the real request's metadata.""" + policy = AutoRouterCompressionPolicy(routing="fake-compress", model=None) + messages = [{"role": "user", "content": "hi"}] + request_kwargs = {"metadata": {}} + await messages_for_routing(policy=policy, messages=messages, request_kwargs=request_kwargs) + assert registered_guardrail.request_data_seen[0] is not request_kwargs + assert request_kwargs == {"metadata": {}} diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index f7fe6ad9d39..c0809e53d2e 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -376,6 +376,60 @@ class TestProxyBaseLLMRequestProcessing: assert "litellm_logging_obj" not in persisted_body json.dumps(persisted_body) + @pytest.mark.asyncio + async def test_common_processing_pre_call_logic_arms_auto_router_compression_before_guardrails( + self, monkeypatch + ): + """arm_pre_call must run before pre_call_hook: an auto router's own compression + policy has to be in `data["metadata"]` (naming the model-side guardrail so it + runs even if it isn't default_on) by the time guardrails see the request.""" + processing_obj = ProxyBaseLLMRequestProcessing(data={}) + mock_request = MagicMock(spec=Request) + mock_request.headers = {} + + async def mock_add_litellm_data_to_request(*args, **kwargs): + return {"model": "smart-router", "messages": [{"role": "user", "content": "hi"}]} + + seen_metadata: dict = {} + + async def mock_pre_call_hook(user_api_key_dict, data, call_type): + seen_metadata.update(data.get("metadata") or {}) + return data + + mock_proxy_logging_obj = MagicMock(spec=ProxyLogging) + mock_proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=mock_pre_call_hook) + monkeypatch.setattr( + litellm.proxy.common_request_processing, + "add_litellm_data_to_request", + mock_add_litellm_data_to_request, + ) + + fake_llm_router = MagicMock() + fake_llm_router.get_model_list.return_value = [ + { + "model_name": "smart-router", + "litellm_params": { + "model": "auto_router/complexity_router", + "auto_router_routing_compression": "none", + "auto_router_model_compression": "headroom-model", + }, + } + ] + mock_proxy_config = MagicMock(spec=ProxyConfig) + mock_proxy_config._get_hierarchical_router_settings = AsyncMock(return_value=None) + + await processing_obj.common_processing_pre_call_logic( + request=mock_request, + general_settings={}, + user_api_key_dict=ProxyUserAPIKeyAuth(api_key="sk-test"), + proxy_logging_obj=mock_proxy_logging_obj, + proxy_config=mock_proxy_config, + route_type="acompletion", + llm_router=fake_llm_router, + ) + + assert seen_metadata.get("guardrails") == ["headroom-model"] + def test_add_dd_apm_tags_for_litellm_call_id_uses_dd_tracing_helper(self, monkeypatch): mock_set_active_span_tag = MagicMock(return_value=True) import litellm.proxy.dd_span_tagger diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 5fc96bcfbb1..cdae3b131ae 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -18,6 +18,7 @@ import pytest import litellm from litellm import Router from litellm.exceptions import MidStreamFallbackError +from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.integrations.custom_logger import CustomLogger from litellm.llms.bedrock.common_utils import BedrockError from litellm.llms.anthropic.experimental_pass_through.messages.agentic_streaming_iterator import ( @@ -9995,6 +9996,156 @@ class TestModelGroupAliasReachesPreRoutingStrategies: ) +class TestAutoRouterCompressionDecoupling: + """An auto router's `auto_router_routing_compression` / `auto_router_model_compression` + decouple what the routing decision sees from what the model call sees. The one + assertion that must hold under any mutation: the strategy can be routed on + compressed text while the caller's own `messages` list - the one that would reach + the model - is never touched.""" + + class _RecordingStrategy: + """Echoes back whatever `messages` it was handed, like every real strategy does.""" + + def __init__(self): + self.received_messages: list[dict] | None = None + + async def async_pre_routing_hook( + self, model, request_kwargs, messages=None, input=None, specific_deployment=False + ): + from litellm.types.router import PreRoutingHookResponse + + self.received_messages = messages + return PreRoutingHookResponse(model="gemini-flash", messages=messages) + + class _CompressingGuardrail(CustomGuardrail): + def __init__(self, guardrail_name: str): + super().__init__(guardrail_name=guardrail_name) + self.call_count = 0 + + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + self.call_count += 1 + structured_messages = inputs.get("structured_messages") or [] + compressed = [ + {**m, "content": f"[COMPRESSED] {m.get('content')}"} for m in structured_messages + ] + return {**inputs, "structured_messages": compressed} + + @staticmethod + def _messages() -> list[dict[str, str]]: + return [{"role": "user", "content": "What is the capital of France?"}] + + def _router(self, marker_litellm_params: dict) -> tuple[litellm.Router, "_RecordingStrategy"]: + from litellm.types.router import TaggedPreRoutingStrategy + + tiers = dict.fromkeys(("SIMPLE", "MEDIUM", "COMPLEX", "REASONING"), "gemini-flash") + router = litellm.Router( + model_list=[ + { + "model_name": "smart-router", + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": {"tiers": tiers}, + "complexity_router_default_model": "gemini-flash", + **marker_litellm_params, + }, + }, + { + "model_name": "gemini-flash", + "litellm_params": {"model": "gemini/gemini-3.6-flash", "mock_response": "routed by the tier"}, + }, + ], + ) + for name in ("auto_routers", "complexity_routers", "adaptive_routers", "quality_routers"): + setattr(router, name, {}) + strategy = self._RecordingStrategy() + router.complexity_routers = {"smart-router": [TaggedPreRoutingStrategy(tags=(), strategy=strategy)]} + return router, strategy + + @pytest.fixture + def registered_guardrail(self): + guardrail = self._CompressingGuardrail(guardrail_name="fake-compress") + litellm.logging_callback_manager.add_litellm_callback(guardrail) + yield guardrail + litellm.logging_callback_manager.remove_callback_from_all_lists(guardrail) + + @pytest.mark.asyncio + async def test_routing_side_compression_never_reaches_the_caller_messages(self, registered_guardrail): + router, strategy = self._router( + { + "auto_router_routing_compression": "fake-compress", + "auto_router_model_compression": "none", + } + ) + original_messages = self._messages() + + response = await router.async_pre_routing_hook( + model="smart-router", request_kwargs={"metadata": {}}, messages=original_messages + ) + + assert strategy.received_messages == [ + {"role": "user", "content": "[COMPRESSED] What is the capital of France?"} + ] + assert response.messages == original_messages + + @pytest.mark.asyncio + async def test_model_side_compression_alone_leaves_routing_uncompressed(self, registered_guardrail): + router, strategy = self._router( + { + "auto_router_routing_compression": "none", + "auto_router_model_compression": "fake-compress", + } + ) + original_messages = self._messages() + + response = await router.async_pre_routing_hook( + model="smart-router", request_kwargs={"metadata": {}}, messages=original_messages + ) + + assert strategy.received_messages == original_messages + assert response.messages == original_messages + assert registered_guardrail.call_count == 0 + + @pytest.mark.asyncio + async def test_same_compression_on_both_hops_compresses_once(self, registered_guardrail): + """The same/different distinction exists so a shared choice does not pay for + compression twice: by the time the router runs, `messages` already reflects + whatever the ordinary pre-call guardrail pipeline did for the model call, so + the routing decision must reuse it rather than calling the guardrail again.""" + router, strategy = self._router( + { + "auto_router_routing_compression": "fake-compress", + "auto_router_model_compression": "fake-compress", + } + ) + # Stands in for what the proxy's ordinary pre-call guardrail pipeline would + # have already produced for the model call, since `auto_router_model_compression` + # names a guardrail: the router never triggers that pipeline itself. + already_compressed_messages = [ + {"role": "user", "content": "[COMPRESSED] What is the capital of France?"} + ] + + response = await router.async_pre_routing_hook( + model="smart-router", request_kwargs={"metadata": {}}, messages=already_compressed_messages + ) + + assert strategy.received_messages == already_compressed_messages + assert response.messages == already_compressed_messages + assert registered_guardrail.call_count == 0 + + @pytest.mark.asyncio + async def test_no_policy_is_fully_unaffected(self, registered_guardrail): + router, strategy = self._router({}) + original_messages = self._messages() + + response = await router.async_pre_routing_hook( + model="smart-router", request_kwargs={"metadata": {}}, messages=original_messages + ) + + assert strategy.received_messages is original_messages + assert response.messages == original_messages + assert registered_guardrail.call_count == 0 + + @pytest.mark.usefixtures("local_model_cost_map") class TestAzureBaseModelFallbackLogging: """When an azure deployment has no base_model but its model name is a known diff --git a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx index 2a024ab7fdf..42115265034 100644 --- a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx @@ -48,6 +48,8 @@ import EscalationKeywords from "./EscalationKeywords"; import KeywordTierRules, { KeywordTierRule } from "./KeywordTierRules"; import SemanticKeywordMatching from "./SemanticKeywordMatching"; import { type DimensionWeights, type TierBoundaries, type TokenThresholds } from "./heuristic_scoring_knobs"; +import CompressionControls from "./CompressionControls"; +import { type AutoRouterCompressionState, DEFAULT_AUTO_ROUTER_COMPRESSION } from "./buildAutoRouterCompression"; export type { DimensionWeights, TierBoundaries, TokenThresholds }; export type { CustomTierSet, TierRow } from "./tier_rows"; @@ -490,6 +492,10 @@ interface ComplexityRouterConfigProps { onMatchThresholdChange?: (threshold: number) => void; escalationKeywords?: string[]; onEscalationKeywordsChange?: (keywords: string[]) => void; + // Optional: not part of complexity_router_config, since it applies to every + // pre-routing strategy, not just the complexity router. + autoRouterCompression?: AutoRouterCompressionState; + onAutoRouterCompressionChange?: (state: AutoRouterCompressionState) => void; showValidationErrors?: boolean; } @@ -611,6 +617,8 @@ const ComplexityRouterConfig: React.FC = ({ onMatchThresholdChange = () => {}, escalationKeywords = [], onEscalationKeywordsChange, + autoRouterCompression = DEFAULT_AUTO_ROUTER_COMPRESSION, + onAutoRouterCompressionChange, showValidationErrors = false, }) => { const customTierSet = value.custom_tier_set; @@ -875,6 +883,32 @@ const ComplexityRouterConfig: React.FC = ({ }, ] : []), + ...(onAutoRouterCompressionChange + ? [ + { + key: "compression", + label: Advanced: Compression, + children: ( + + onAutoRouterCompressionChange({ + ...autoRouterCompression, + routing, + sameAsRouting: routing === undefined ? true : autoRouterCompression.sameAsRouting, + }) + } + sameAsRouting={autoRouterCompression.sameAsRouting} + onSameAsRoutingChange={(sameAsRouting) => + onAutoRouterCompressionChange({ ...autoRouterCompression, sameAsRouting }) + } + model={autoRouterCompression.model} + onModelChange={(model) => onAutoRouterCompressionChange({ ...autoRouterCompression, model })} + /> + ), + }, + ] + : []), ...(onKeywordTierRulesChange || onSemanticMatchingEnabledChange ? [ { diff --git a/ui/litellm-dashboard/src/components/add_model/CompressionControls.tsx b/ui/litellm-dashboard/src/components/add_model/CompressionControls.tsx new file mode 100644 index 00000000000..a0a240f76b0 --- /dev/null +++ b/ui/litellm-dashboard/src/components/add_model/CompressionControls.tsx @@ -0,0 +1,93 @@ +import { SimpleTooltip } from "@/components/ui/tooltip"; +import { SearchSelect, SearchSelectOption } from "@/components/shared/SearchSelect"; +import { Label } from "@/components/ui/label"; +import { RadioGroup, RadioGroupItem } from "@/components/ui/radio-group"; +import { Info } from "lucide-react"; +import React from "react"; +import { useGuardrails } from "@/app/(dashboard)/hooks/guardrails/useGuardrails"; +import { COMPRESSION_GUARDRAIL_PROVIDER } from "@/app/(dashboard)/cost-optimization/_components/helpers"; +import { NO_COMPRESSION } from "./buildAutoRouterCompression"; + +interface CompressionControlsProps { + routing: string | undefined; + onRoutingChange: (value: string | undefined) => void; + sameAsRouting: boolean; + onSameAsRoutingChange: (same: boolean) => void; + model: string | undefined; + onModelChange: (value: string | undefined) => void; +} + +const NONE_OPTION: SearchSelectOption = { label: "None (no compression)", value: NO_COMPRESSION }; + +const CompressionControls: React.FC = ({ + routing, + onRoutingChange, + sameAsRouting, + onSameAsRoutingChange, + model, + onModelChange, +}) => { + const { data } = useGuardrails(); + const compressionOptions: SearchSelectOption[] = (data?.guardrails ?? []) + .filter((g) => (g.litellm_params?.guardrail ?? "").toString().toLowerCase() === COMPRESSION_GUARDRAIL_PROVIDER) + .map((g) => ({ label: g.guardrail_name, value: g.guardrail_name })); + const options: SearchSelectOption[] = [NONE_OPTION, ...compressionOptions]; + + return ( +
+
+
+ Routing decision + + + +
+ onRoutingChange(value === "" ? undefined : value)} + placeholder="Inherit from the request's own compression guardrails" + emptyText="No compression guardrails found" + aria-label="Routing decision compression" + /> +
+ + {routing !== undefined && ( +
+ Model call + onSameAsRoutingChange(value === "same")} + className="w-full" + > +
+ + +
+ + + {!sameAsRouting && ( +
+ onModelChange(value === "" ? undefined : value)} + placeholder="None (no compression)" + emptyText="No compression guardrails found" + aria-label="Model call compression" + /> +
+ )} +
+ )} +
+ ); +}; + +export default CompressionControls; diff --git a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx index d2f6b10c3a6..5605a993ded 100644 --- a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx @@ -1,4 +1,12 @@ -import { renderWithProviders, screen, waitFor, within, fireEvent, testQueryClient } from "../../../tests/test-utils"; +import { + renderWithProviders, + screen, + waitFor, + within, + fireEvent, + testQueryClient, + chooseSelectOption, +} from "../../../tests/test-utils"; import userEvent from "@testing-library/user-event"; import { vi } from "vitest"; import AddAutoRouterTab from "./add_auto_router_tab"; @@ -522,6 +530,71 @@ describe("AddAutoRouterTab", () => { ); }); + describe("prompt compression", () => { + it("leaves both compression keys out of the create payload when the section is untouched", async () => { + const user = userEvent.setup(); + vi.mocked(getMissingTiersError).mockReturnValue(null); + + renderWithProviders(); + + await user.type(screen.getByPlaceholderText(/smart_router/i), "no-compression-router"); + await user.click(screen.getByRole("button", { name: /add auto router/i })); + + await waitFor(() => expect(handleAddAutoRouterSubmit).toHaveBeenCalled()); + const submitted = vi.mocked(handleAddAutoRouterSubmit).mock.calls.at(-1)?.[0]; + expect(submitted).not.toHaveProperty("auto_router_routing_compression"); + expect(submitted).not.toHaveProperty("auto_router_model_compression"); + }); + + it("mirrors an explicit no-compression routing choice onto the model call by default", async () => { + const user = userEvent.setup(); + vi.mocked(getMissingTiersError).mockReturnValue(null); + + renderWithProviders(); + + await user.type(screen.getByPlaceholderText(/smart_router/i), "no-compression-explicit-router"); + expandDetailedConfiguration(); + await user.click(screen.getByText("Advanced: Compression")); + await chooseSelectOption( + user, + screen.getByRole("combobox", { name: "Routing decision compression" }), + "None (no compression)", + ); + + await user.click(screen.getByRole("button", { name: /add auto router/i })); + + await waitFor(() => expect(handleAddAutoRouterSubmit).toHaveBeenCalled()); + const submitted = vi.mocked(handleAddAutoRouterSubmit).mock.calls.at(-1)?.[0]; + expect(submitted?.auto_router_routing_compression).toBe("none"); + expect(submitted?.auto_router_model_compression).toBe("none"); + }); + + it("defaults the model call to none when different is chosen but nothing is picked there", async () => { + const user = userEvent.setup(); + vi.mocked(getMissingTiersError).mockReturnValue(null); + + renderWithProviders(); + + await user.type(screen.getByPlaceholderText(/smart_router/i), "different-compression-router"); + expandDetailedConfiguration(); + await user.click(screen.getByText("Advanced: Compression")); + await chooseSelectOption( + user, + screen.getByRole("combobox", { name: "Routing decision compression" }), + "None (no compression)", + ); + await user.click(screen.getByText("Use a different compression")); + expect(screen.getByRole("combobox", { name: "Model call compression" })).toBeInTheDocument(); + + await user.click(screen.getByRole("button", { name: /add auto router/i })); + + await waitFor(() => expect(handleAddAutoRouterSubmit).toHaveBeenCalled()); + const submitted = vi.mocked(handleAddAutoRouterSubmit).mock.calls.at(-1)?.[0]; + expect(submitted?.auto_router_routing_compression).toBe("none"); + expect(submitted?.auto_router_model_compression).toBe("none"); + }); + }); + // The scalar floor is the one scorer knob with no group dict behind it, so its wiring into the create // payload is only proven end to end. 0 is the case a truthy check would silently drop. it("carries a reasoning override floor of 0 through to the create payload", async () => { diff --git a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx index a548c2c6533..decacac6501 100644 --- a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx +++ b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx @@ -32,6 +32,11 @@ import ComplexityRouterConfig, { } from "./ComplexityRouterConfig"; import { KeywordTierRule } from "./KeywordTierRules"; import { DEFAULT_ESCALATION_KEYWORDS } from "./EscalationKeywords"; +import { + type AutoRouterCompressionState, + buildAutoRouterCompressionParams, + DEFAULT_AUTO_ROUTER_COMPRESSION, +} from "./buildAutoRouterCompression"; import { DEFAULT_MATCH_THRESHOLD } from "./SemanticKeywordMatching"; import { BuildComplexityRouterConfigParams, @@ -194,6 +199,9 @@ const AddAutoRouterTab: React.FC = ({ const [embeddingModel, setEmbeddingModel] = useState(undefined); const [matchThreshold, setMatchThreshold] = useState(DEFAULT_MATCH_THRESHOLD); const [escalationKeywords, setEscalationKeywords] = useState(DEFAULT_ESCALATION_KEYWORDS); + const [autoRouterCompression, setAutoRouterCompression] = useState( + DEFAULT_AUTO_ROUTER_COMPRESSION, + ); const [showValidationErrors, setShowValidationErrors] = useState(false); const [editingTiers, setEditingTiers] = useState(false); const [isSubmitting, setIsSubmitting] = useState(false); @@ -461,6 +469,7 @@ const AddAutoRouterTab: React.FC = ({ model_type: "complexity_router", complexity_router_config: complexityRouterConfigPayload, model_access_group: form.getValues("model_access_group"), + ...buildAutoRouterCompressionParams(autoRouterCompression), }; await handleAddAutoRouterSubmit(submitValues, accessToken, () => form.reset(EMPTY_FORM_VALUES), handleOk); @@ -666,6 +675,8 @@ const AddAutoRouterTab: React.FC = ({ onMatchThresholdChange={setMatchThreshold} escalationKeywords={escalationKeywords} onEscalationKeywordsChange={setEscalationKeywords} + autoRouterCompression={autoRouterCompression} + onAutoRouterCompressionChange={setAutoRouterCompression} showValidationErrors={showValidationErrors} /> diff --git a/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.test.ts b/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.test.ts new file mode 100644 index 00000000000..b917fcedaa2 --- /dev/null +++ b/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.test.ts @@ -0,0 +1,93 @@ +import { + buildAutoRouterCompressionParams, + DEFAULT_AUTO_ROUTER_COMPRESSION, + hydrateAutoRouterCompression, + NO_COMPRESSION, +} from "./buildAutoRouterCompression"; + +describe("buildAutoRouterCompressionParams", () => { + it("omits both keys when routing was never configured", () => { + expect(buildAutoRouterCompressionParams(DEFAULT_AUTO_ROUTER_COMPRESSION)).toEqual({}); + }); + + it("mirrors routing onto model when same-as-routing is chosen", () => { + const params = buildAutoRouterCompressionParams({ + routing: "headroom-a", + sameAsRouting: true, + model: undefined, + }); + expect(params).toEqual({ + auto_router_routing_compression: "headroom-a", + auto_router_model_compression: "headroom-a", + }); + }); + + it("uses the explicit model choice when different is chosen", () => { + const params = buildAutoRouterCompressionParams({ + routing: "headroom-a", + sameAsRouting: false, + model: "headroom-b", + }); + expect(params).toEqual({ + auto_router_routing_compression: "headroom-a", + auto_router_model_compression: "headroom-b", + }); + }); + + it("defaults the model side to none when different is chosen but nothing is picked", () => { + const params = buildAutoRouterCompressionParams({ + routing: "headroom-a", + sameAsRouting: false, + model: undefined, + }); + expect(params).toEqual({ + auto_router_routing_compression: "headroom-a", + auto_router_model_compression: NO_COMPRESSION, + }); + }); + + it("sends the none sentinel when routing itself is explicitly turned off", () => { + const params = buildAutoRouterCompressionParams({ + routing: NO_COMPRESSION, + sameAsRouting: true, + model: undefined, + }); + expect(params).toEqual({ + auto_router_routing_compression: NO_COMPRESSION, + auto_router_model_compression: NO_COMPRESSION, + }); + }); +}); + +describe("hydrateAutoRouterCompression", () => { + it("returns the default state when neither key is set", () => { + expect(hydrateAutoRouterCompression({})).toEqual(DEFAULT_AUTO_ROUTER_COMPRESSION); + }); + + it("is same-as-routing when the model value matches routing", () => { + const state = hydrateAutoRouterCompression({ + auto_router_routing_compression: "headroom-a", + auto_router_model_compression: "headroom-a", + }); + expect(state).toEqual({ routing: "headroom-a", sameAsRouting: true, model: undefined }); + }); + + it("is different when the model value diverges from routing", () => { + const state = hydrateAutoRouterCompression({ + auto_router_routing_compression: "headroom-a", + auto_router_model_compression: "headroom-b", + }); + expect(state).toEqual({ routing: "headroom-a", sameAsRouting: false, model: "headroom-b" }); + }); + + it("treats a missing model key as same-as-routing", () => { + const state = hydrateAutoRouterCompression({ auto_router_routing_compression: "headroom-a" }); + expect(state).toEqual({ routing: "headroom-a", sameAsRouting: true, model: undefined }); + }); + + it("round-trips through buildAutoRouterCompressionParams", () => { + const original = { auto_router_routing_compression: "headroom-a", auto_router_model_compression: "none" }; + const rebuilt = buildAutoRouterCompressionParams(hydrateAutoRouterCompression(original)); + expect(rebuilt).toEqual(original); + }); +}); diff --git a/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.ts b/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.ts new file mode 100644 index 00000000000..49180d5b4ce --- /dev/null +++ b/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.ts @@ -0,0 +1,52 @@ +/** + * Maps the auto router's compression form state to the two flat litellm_params keys + * the backend reads (litellm.proxy.guardrails.auto_router_compression), and back. + * + * `routing` being undefined means the section was never touched: both keys are + * omitted from the payload, and the request's own compression guardrails apply to + * both hops unchanged. Once `routing` has a value (a guardrail name, or the "none" + * sentinel for explicit no-compression), the auto router is authoritative and the + * model side always gets a concrete value too, mirroring `routing` when same-as + * is chosen and defaulting to "none" otherwise. + */ + +export const NO_COMPRESSION = "none"; + +export interface AutoRouterCompressionState { + routing: string | undefined; + sameAsRouting: boolean; + model: string | undefined; +} + +export interface AutoRouterCompressionLitellmParams { + auto_router_routing_compression?: string; + auto_router_model_compression?: string; +} + +export const DEFAULT_AUTO_ROUTER_COMPRESSION: AutoRouterCompressionState = { + routing: undefined, + sameAsRouting: true, + model: undefined, +}; + +export const buildAutoRouterCompressionParams = ( + state: AutoRouterCompressionState, +): AutoRouterCompressionLitellmParams => { + if (state.routing === undefined) return {}; + return { + auto_router_routing_compression: state.routing, + auto_router_model_compression: state.sameAsRouting ? state.routing : (state.model ?? NO_COMPRESSION), + }; +}; + +export const hydrateAutoRouterCompression = (litellmParams: { + auto_router_routing_compression?: string | null; + auto_router_model_compression?: string | null; +}): AutoRouterCompressionState => { + const routing = litellmParams.auto_router_routing_compression ?? undefined; + if (routing === undefined) return DEFAULT_AUTO_ROUTER_COMPRESSION; + + const model = litellmParams.auto_router_model_compression ?? undefined; + const sameAsRouting = model === undefined || model === routing; + return { routing, sameAsRouting, model: sameAsRouting ? undefined : model }; +}; diff --git a/ui/litellm-dashboard/src/components/add_model/handle_add_auto_router_submit.tsx b/ui/litellm-dashboard/src/components/add_model/handle_add_auto_router_submit.tsx index 9385836ce1a..59d9ecf205e 100644 --- a/ui/litellm-dashboard/src/components/add_model/handle_add_auto_router_submit.tsx +++ b/ui/litellm-dashboard/src/components/add_model/handle_add_auto_router_submit.tsx @@ -1,8 +1,9 @@ import { modelCreateCall } from "../networking"; import { toast } from "@/lib/toast"; import type { ComplexityRouterConfigPayload } from "./build_complexity_router_config"; +import type { AutoRouterCompressionLitellmParams } from "./buildAutoRouterCompression"; -export interface AddAutoRouterValues { +export interface AddAutoRouterValues extends AutoRouterCompressionLitellmParams { auto_router_name: string; auto_router_default_model: string | undefined; model_type: "complexity_router"; @@ -24,6 +25,8 @@ export const handleAddAutoRouterSubmit = async ( model: "auto_router/complexity_router", complexity_router_config: values.complexity_router_config, complexity_router_default_model: values.auto_router_default_model, + auto_router_routing_compression: values.auto_router_routing_compression, + auto_router_model_compression: values.auto_router_model_compression, }, model_info: { ...(values.team_id ? { team_id: values.team_id } : {}), 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.test.tsx index 970bcaa545f..b93db8d963e 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.test.tsx @@ -1029,3 +1029,84 @@ describe("EditAutoRouterModal with a stored custom tier set", () => { expect(savedConfig().tier_model_configs).toEqual(CUSTOM_STORED.tier_model_configs); }); }); + +describe("EditAutoRouterModal prompt compression", () => { + beforeEach(() => { + modelPatchUpdateCall.mockClear(); + }); + + const savedLitellmParams = () => { + const [, payload] = modelPatchUpdateCall.mock.calls.at(-1) ?? []; + return payload?.litellm_params; + }; + + const renderWithStoredCompression = ( + compression?: { auto_router_routing_compression?: string; auto_router_model_compression?: string }, + ) => + renderWithProviders( + , + ); + + it("leaves both compression keys out of an untouched save when none were stored", async () => { + const user = userEvent.setup(); + renderWithStoredCompression(); + + await user.click(await screen.findByRole("button", { name: /save changes/i })); + + await waitFor(() => expect(modelPatchUpdateCall).toHaveBeenCalled()); + expect(savedLitellmParams()).not.toHaveProperty("auto_router_routing_compression"); + expect(savedLitellmParams()).not.toHaveProperty("auto_router_model_compression"); + }); + + it("preserves a stored same-as-routing compression through an untouched open-and-save", async () => { + const user = userEvent.setup(); + renderWithStoredCompression({ + auto_router_routing_compression: "headroom-a", + auto_router_model_compression: "headroom-a", + }); + + 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"); + }); + + it("shows a stored different-compression choice as Use a different compression, not Same", async () => { + const user = userEvent.setup(); + renderWithStoredCompression({ + auto_router_routing_compression: "headroom-a", + auto_router_model_compression: "none", + }); + + await user.click(await screen.findByText("Advanced: Compression")); + + expect(await screen.findByRole("combobox", { name: "Routing decision compression" })).toHaveValue("headroom-a"); + expect(screen.getByRole("radio", { name: "Use a different compression" })).toBeChecked(); + expect(screen.getByRole("combobox", { name: "Model call compression" })).toHaveValue("None (no compression)"); + }); + + it("preserves a stored different-compression choice through an untouched open-and-save", async () => { + const user = userEvent.setup(); + renderWithStoredCompression({ + auto_router_routing_compression: "headroom-a", + auto_router_model_compression: "none", + }); + + 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("none"); + }); +}); 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 ea1e5cba6a3..a4852f9e784 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 @@ -41,6 +41,12 @@ import { } from "../add_model/build_complexity_router_config"; import { KeywordTierRule } from "../add_model/KeywordTierRules"; import { DEFAULT_MATCH_THRESHOLD } from "../add_model/SemanticKeywordMatching"; +import { + type AutoRouterCompressionState, + buildAutoRouterCompressionParams, + DEFAULT_AUTO_ROUTER_COMPRESSION, + hydrateAutoRouterCompression, +} from "../add_model/buildAutoRouterCompression"; import { hydrateKeywordTierRules } from "../add_model/complexity_router_keywords"; import { hydrateDimensionWeights, @@ -424,6 +430,9 @@ const EditAutoRouterModal: React.FC = ({ const [semanticMatchingEnabled, setSemanticMatchingEnabled] = useState(false); const [embeddingModel, setEmbeddingModel] = useState(undefined); const [matchThreshold, setMatchThreshold] = useState(DEFAULT_MATCH_THRESHOLD); + const [autoRouterCompression, setAutoRouterCompression] = useState( + DEFAULT_AUTO_ROUTER_COMPRESSION, + ); const [complexityRouterConfig, setComplexityRouterConfig] = useState({ tiers: { SIMPLE: [], MEDIUM: [], COMPLEX: [], REASONING: [] }, classifier_type: "heuristic", @@ -516,6 +525,12 @@ const EditAutoRouterModal: React.FC = ({ setMatchThreshold( typeof parsedConfig.match_threshold === "number" ? parsedConfig.match_threshold : DEFAULT_MATCH_THRESHOLD, ); + setAutoRouterCompression( + hydrateAutoRouterCompression({ + auto_router_routing_compression: modelData.litellm_params?.auto_router_routing_compression, + auto_router_model_compression: modelData.litellm_params?.auto_router_model_compression, + }), + ); form.reset({ ...EMPTY_FORM_VALUES, @@ -628,6 +643,7 @@ const EditAutoRouterModal: React.FC = ({ ...modelData.litellm_params, complexity_router_config: updatedConfig, complexity_router_default_model: defaultModel, + ...buildAutoRouterCompressionParams(autoRouterCompression), }; const updatedModelInfo = { ...modelData.model_info, @@ -749,6 +765,8 @@ const EditAutoRouterModal: React.FC = ({ onMatchThresholdChange={setMatchThreshold} escalationKeywords={escalationKeywords} onEscalationKeywordsChange={setEscalationKeywords} + autoRouterCompression={autoRouterCompression} + onAutoRouterCompressionChange={setAutoRouterCompression} /> ) : ( diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 5f3cca49644..8f6a0700517 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -29222,6 +29222,10 @@ export interface components { auto_router_embedding_model?: string | null; /** Auto Router Max Input Chars */ auto_router_max_input_chars?: number | null; + /** Auto Router Model Compression */ + auto_router_model_compression?: string | null; + /** Auto Router Routing Compression */ + auto_router_routing_compression?: string | null; /** Aws Access Key Id */ aws_access_key_id?: string | null; /** Aws Batch Role Arn */ @@ -39275,6 +39279,10 @@ export interface components { auto_router_embedding_model?: string | null; /** Auto Router Max Input Chars */ auto_router_max_input_chars?: number | null; + /** Auto Router Model Compression */ + auto_router_model_compression?: string | null; + /** Auto Router Routing Compression */ + auto_router_routing_compression?: string | null; /** Aws Access Key Id */ aws_access_key_id?: string | null; /** Aws Batch Role Arn */