diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index c462b1edb98..2d66a280663 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -948,6 +948,23 @@ class CustomGuardrail(CustomLogger): """ return False + def _suppressed_by_auto_router_compression(self) -> bool: + """True when an auto router's own compression policy suppresses this guardrail. + + Reads request-scoped state set by `arm_pre_call`, never request metadata. The + caller controls metadata, and metadata reaches spend logs the caller can read, + so a suppression list carried there would be one a request could replay to + switch off a PII or content-filter guardrail for itself. + """ + name: Final = self.guardrail_name + if not name: + return False + from litellm.proxy.guardrails.auto_router_compression import ( + suppressed_compression_guardrails, + ) + + return name in suppressed_compression_guardrails() + def should_run_guardrail( self, data, @@ -956,6 +973,9 @@ class CustomGuardrail(CustomLogger): """ Returns True if the guardrail should be run on the event_type """ + if self._suppressed_by_auto_router_compression(): + 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..5e6c9b34332 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. + 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..98707e7ddca --- /dev/null +++ b/litellm/proxy/guardrails/auto_router_compression.py @@ -0,0 +1,267 @@ +""" +Decouples prompt compression between an auto router's routing decision and the model +it routes to, via ``auto_router_routing_compression`` / ``auto_router_model_compression`` +on the marker deployment: a guardrail name, or ``"none"``. + +Neither key set inherits today's behaviour. Either key set makes the auto router +authoritative and suppresses every other compression guardrail for that request. +""" + +import contextvars +from collections.abc import Iterable, Mapping, Sequence +from dataclasses import dataclass +from typing import TYPE_CHECKING, Final + +from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.core_helpers import 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 + +COMPRESSION_GUARDRAIL_PROVIDERS: Final = frozenset({"headroom", "compresr"}) +_NO_COMPRESSION: Final = "none" + +# A ContextVar, not metadata: metadata reaches spend logs the caller can read, and a +# suppression list they can read is one they can replay to disable any guardrail. +_suppressed_compression_guardrails: Final[contextvars.ContextVar[frozenset[str]]] = contextvars.ContextVar( + "litellm_auto_router_suppressed_compression_guardrails", default=frozenset() +) + + +def suppressed_compression_guardrails() -> frozenset[str]: + """Names of the compression guardrails this request's auto router suppresses.""" + return _suppressed_compression_guardrails.get() + + +# Only the proxy calls `arm_pre_call`, so on the SDK path nothing arms and nothing +# compresses; the router must not assume the model hop already ran. +_model_hop_armed: Final[contextvars.ContextVar[bool]] = contextvars.ContextVar( + "litellm_auto_router_model_hop_armed", default=False +) + + +def model_hop_compression_armed() -> bool: + """True when this request's model-side compression guardrail was actually armed.""" + return _model_hop_armed.get() + + +@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, + request_tags: Sequence[str], +) -> AutoRouterCompressionPolicy | None: + """The compression policy of the auto router marker `model_alias` resolves to. + + Pre-call arming and the routing hook both resolve through here, so an alias with + several tag-scoped markers cannot suppress under one and then route under another. + """ + if llm_router is None: + return None + deployments: Final = llm_router.get_model_list(model_name=model_alias, team_id=team_id) or () + markers: Final = tuple( + litellm_params + for deployment in deployments + if isinstance(litellm_params := deployment.get("litellm_params"), Mapping) # pyright: ignore[reportUnnecessaryIsInstance] # filters out non-Mapping + and str(litellm_params.get("model", "")).startswith(AUTO_ROUTER_MODEL_PREFIX) + ) + requested: Final = frozenset(request_tags) + tag_matched: Final = tuple( + params for params in markers if (tags := params.get("tags")) and requested.issuperset(frozenset(tags)) + ) + # Untagged only: a marker scoped to tags this request lacks describes other traffic. + untagged: Final = tuple(params for params in markers if not params.get("tags")) + # Lazy, so the first marker carrying a policy wins and the rest are never read. + candidates: Final = (policy_from_litellm_params(params) for params in (*tag_matched, *untagged)) + return next((policy for policy in candidates if policy is not None), None) + + +def team_id_from_request(request_kwargs: Mapping[str, object]) -> str | None: + """The caller's team id, from whichever metadata bucket this surface writes to.""" + for meta_key in ("metadata", "litellm_metadata"): + meta = request_kwargs.get(meta_key) + if isinstance(meta, Mapping): + team_id = meta.get("user_api_key_team_id") + if isinstance(team_id, str): + return team_id + return None + + +def _compression_guardrail_classes() -> tuple[type, ...]: + """The registered guardrail classes whose provider compresses prompts.""" + from litellm.proxy.guardrails.guardrail_registry import guardrail_class_registry + + return tuple(cls for name, cls in guardrail_class_registry.items() if name in COMPRESSION_GUARDRAIL_PROVIDERS) + + +def is_compression_guardrail(guardrail: object) -> bool: + """Whether `guardrail` is an instance of a compression guardrail provider. + + Both hops validate through here: the policy fields are operator-supplied names, and + an unvalidated one would get handed the conversation and invoked. + """ + classes: Final = _compression_guardrail_classes() + return bool(classes) and isinstance(guardrail, classes) + + +def _active_compression_guardrails() -> tuple["CustomGuardrail", ...]: + """Every currently-active guardrail whose type is a compression guardrail.""" + import litellm + from litellm.integrations.custom_guardrail import CustomGuardrail + + if not _compression_guardrail_classes(): + return () + active: Final = litellm.logging_callback_manager.get_custom_loggers_for_type(callback_type=CustomGuardrail) + return tuple(cb for cb in active if is_compression_guardrail(cb) and cb.guardrail_name) + + +async def arm_pre_call( + data: dict[str, object], # mutable-ok: arms the live request dict in place + llm_router: "Router | None", +) -> None: + """Apply an auto router's compression policy, if any, before guardrails run. + + Suppresses every other compression guardrail and re-enables the model-side + guardrail the policy names (if any) even when it isn't ``default_on``. + """ + _suppressed_compression_guardrails.set(frozenset()) + _model_hop_armed.set(False) + if llm_router is None: + return + + model_alias: Final = data.get("model") + if not isinstance(model_alias, str) or not model_alias: + return + + from litellm.router_strategy.tag_based_routing import ( + _get_tags_from_request_kwargs, # pyright: ignore[reportPrivateUsage] # used in router.py and budget_limiter.py too + ) + + policy: Final = policy_for_model( + llm_router=llm_router, + model_alias=model_alias, + team_id=team_id_from_request(data), + request_tags=_get_tags_from_request_kwargs(data), + ) + if policy is None: + return + + _suppressed_compression_guardrails.set( + frozenset( + name + for guardrail in _active_compression_guardrails() + if (name := guardrail.guardrail_name) and name != policy.model + ) + ) + + # Arming adds the name to `metadata["guardrails"]`, which runs it even if not default_on. + armed_model_hop: Final = policy.model is not None and any( + guardrail.guardrail_name == policy.model for guardrail in _active_compression_guardrails() + ) + if policy.model is not None and not armed_model_hop: + verbose_proxy_logger.warning( + "AutoRouter compression: '%s' is not an active compression guardrail; the model hop is uncompressed", + policy.model, + ) + + if armed_model_hop: + _model_hop_armed.set(True) + _, metadata = get_or_create_metadata_bucket(data) + requested: Final = metadata.get("guardrails") + existing: Final = tuple(requested) if isinstance(requested, (list, tuple)) else () + if policy.model not in existing: + # A list: litellm_pre_call_utils isinstance-checks this key and drops a tuple. + metadata["guardrails"] = [*existing, policy.model] # mutable-ok: this key's contract is a list + + +def _as_routing_messages( + messages: Iterable[Mapping[str, object]], +) -> list[dict[str, object]]: # mutable-ok: shape fixed by the pre-routing hook protocol + """A fresh, independently mutable copy, the shape the pre-routing hook takes.""" + return [dict(message) for message in messages] # mutable-ok: shape fixed by the pre-routing hook protocol + + +async def messages_for_routing( + policy: AutoRouterCompressionPolicy | None, + # list[dict], not Sequence[Mapping]: fixed by the async_pre_routing_hook protocol. + messages: list[dict[str, object]] | None, # mutable-ok: shape fixed by the pre-routing hook protocol + request_kwargs: Mapping[str, object], +) -> list[dict[str, object]] | None: # mutable-ok: shape fixed by the pre-routing hook protocol + """Messages to use for a routing decision, per `policy.routing`. None means the + caller should route on whatever it already has. + + Reads the live messages, never a pre-guardrail copy: this compresses through a real + guardrail that POSTs the text out, so routing on a pre-masking snapshot would leak + what the masking guardrail stripped. When the model hop already compressed and the + hops differ, routing therefore reads the compressed text rather than the original. + """ + if policy is None or policy.routing is None: + return None + + if not messages: + return None + + from litellm.proxy.common_utils.registry_read_through import ( + get_initialized_guardrail_with_read_through, + ) + + 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 _as_routing_messages(messages) + + if not is_compression_guardrail(guardrail): + verbose_proxy_logger.warning( + "AutoRouter compression: guardrail '%s' is not a compression guardrail; routing on uncompressed messages", + policy.routing, + ) + return _as_routing_messages(messages) + + inputs: Final[GenericGuardrailAPIInputs] = { + "structured_messages": _as_routing_messages(messages) # pyright: ignore[reportAssignmentType] # plain dicts, not AllMessageValues; see headroom.py's own use of this shape + } + model: Final = request_kwargs.get("model") + # Throwaway: apply_guardrail writes stats here, so routing never double-counts into + # extract_compression_saved_tokens. + stats_sink: Final = {"messages": messages, "model": model} # mutable-ok: apply_guardrail writes its stats here + result: Final = await guardrail.apply_guardrail( + inputs=inputs, + request_data=stats_sink, + input_type="request", + ) + compressed: Final = result.get("structured_messages") + return compressed if isinstance(compressed, list) else _as_routing_messages(messages) diff --git a/litellm/router.py b/litellm/router.py index 76da3a857df..6d6efad9f42 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -8636,7 +8636,7 @@ class Router: raise ValueError(ptu_error) zeroed_pricing: Final = zeroed_ptu_pricing(_model_info, _litellm_params) if config_sourced else None litellm_params: Final[LiteLLM_Params] = LiteLLM_Params( - **( + **( # pyright: ignore[reportArgumentType] # untyped merged dict; already true for every field here _litellm_params if zeroed_pricing is None else MappingProxyType({**_litellm_params, **zeroed_pricing}) @@ -13033,13 +13033,48 @@ class Router: ) return None - pre_routing_hook_response: Final = await selected_strategy.strategy.async_pre_routing_hook( + from litellm.proxy.guardrails.auto_router_compression import ( + messages_for_routing, + model_hop_compression_armed, + policy_for_model, + team_id_from_request, + ) + + # Same tag-aware lookup the proxy's pre-call arming used, so an alias with + # several tag-scoped markers cannot suppress under one and route under another. + compression_policy: Final = policy_for_model( + llm_router=self, + model_alias=registered_model_name, + team_id=team_id_from_request(request_kwargs), + request_tags=_get_tags_from_request_kwargs(request_kwargs), + ) + # Shared compression already ran in the pre-call hook, so reuse it rather than + # compressing twice. Conditional on arming having actually happened: only the + # proxy arms, and on the SDK path the shortcut would skip both hops entirely. + needs_independent_routing_compression: Final = compression_policy is not None and not ( + compression_policy.is_same and compression_policy.model is not None and model_hop_compression_armed() + ) + routing_messages: Final = ( + await messages_for_routing(policy=compression_policy, messages=messages, request_kwargs=request_kwargs) + if needs_independent_routing_compression + else None + ) + + routed: 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, ) + # Routing-only compression must not leak into the response: the model call and + # deployment-context filtering key off this field. Compared by value, since + # pydantic rebuilds the list rather than keeping the object passed in. + pre_routing_hook_response: Final = ( + routed.model_copy(update={"messages": messages}) # mutable-ok: pydantic's model_copy takes a dict + if routed is not None and routing_messages is not None and routed.messages == routing_messages + else routed + ) self._record_routing_decision( request_kwargs=request_kwargs, routing_decision=(pre_routing_hook_response.routing_decision if pre_routing_hook_response else None), diff --git a/litellm/types/router.py b/litellm/types/router.py index 728d1037f3d..2dad22751de 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -361,6 +361,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 5052cd6ef48..2118fe77aad 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3769,6 +3769,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 49a52157c8a..885bd1d4d72 100644 --- a/tests/test_litellm/integrations/test_custom_guardrail.py +++ b/tests/test_litellm/integrations/test_custom_guardrail.py @@ -17,7 +17,6 @@ if TYPE_CHECKING: class TestCustomGuardrailDeploymentHook: - @pytest.mark.asyncio async def test_async_pre_call_deployment_hook_no_guardrails(self): """Test that method returns kwargs unchanged when no guardrails are present""" @@ -30,18 +29,14 @@ class TestCustomGuardrailDeploymentHook: "guardrails": None, } - result = await custom_guardrail.async_pre_call_deployment_hook( - kwargs=kwargs, call_type=CallTypes.completion - ) + result = await custom_guardrail.async_pre_call_deployment_hook(kwargs=kwargs, call_type=CallTypes.completion) assert result == kwargs # Test with guardrails as non-list kwargs["guardrails"] = "not_a_list" - result = await custom_guardrail.async_pre_call_deployment_hook( - kwargs=kwargs, call_type=CallTypes.completion - ) + result = await custom_guardrail.async_pre_call_deployment_hook(kwargs=kwargs, call_type=CallTypes.completion) assert result == kwargs @@ -68,9 +63,7 @@ class TestCustomGuardrailDeploymentHook: "user_api_key_request_route": "test_route", } - result = await custom_guardrail.async_pre_call_deployment_hook( - kwargs=kwargs, call_type=CallTypes.completion - ) + result = await custom_guardrail.async_pre_call_deployment_hook(kwargs=kwargs, call_type=CallTypes.completion) # Verify async_pre_call_hook was called with correct parameters custom_guardrail.async_pre_call_hook.assert_called_once() @@ -103,9 +96,7 @@ class TestCustomGuardrailDeploymentHook: super().__init__(guardrail_name="g1", default_on=True) self.pre_call_count = 0 - async def async_pre_call_hook( - self, user_api_key_dict, cache, data, call_type - ): + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): self.pre_call_count += 1 return data @@ -118,9 +109,7 @@ class TestCustomGuardrailDeploymentHook: } guardrail.mark_pre_call_hook_ran(kwargs) - await guardrail.async_pre_call_deployment_hook( - kwargs=kwargs, call_type=CallTypes.completion - ) + await guardrail.async_pre_call_deployment_hook(kwargs=kwargs, call_type=CallTypes.completion) assert guardrail.pre_call_count == 0 @@ -134,9 +123,7 @@ class TestCustomGuardrailDeploymentHook: super().__init__(guardrail_name="g1", default_on=True) self.pre_call_count = 0 - async def async_pre_call_hook( - self, user_api_key_dict, cache, data, call_type - ): + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): self.pre_call_count += 1 return data @@ -148,9 +135,7 @@ class TestCustomGuardrailDeploymentHook: "metadata": {}, } - await guardrail.async_pre_call_deployment_hook( - kwargs=kwargs, call_type=CallTypes.completion - ) + await guardrail.async_pre_call_deployment_hook(kwargs=kwargs, call_type=CallTypes.completion) assert guardrail.pre_call_count == 1 @@ -179,9 +164,7 @@ class TestCustomGuardrailDeploymentHook: super().__init__(guardrail_name="g1", default_on=True) self.pre_call_count = 0 - async def async_pre_call_hook( - self, user_api_key_dict, cache, data, call_type - ): + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): self.pre_call_count += 1 return data @@ -193,15 +176,12 @@ class TestCustomGuardrailDeploymentHook: "metadata": {PRE_CALL_EXECUTED_GUARDRAILS_KEY: ["g1"]}, } - await guardrail.async_pre_call_deployment_hook( - kwargs=kwargs, call_type=CallTypes.completion - ) + await guardrail.async_pre_call_deployment_hook(kwargs=kwargs, call_type=CallTypes.completion) assert guardrail.pre_call_count == 1 class TestCustomGuardrailShouldRunGuardrail: - def test_should_run_guardrail_with_litellm_metadata(self): """Test that should_run_guardrail works with litellm_metadata pattern""" from litellm.types.guardrails import GuardrailEventHooks @@ -218,9 +198,7 @@ class TestCustomGuardrailShouldRunGuardrail: "litellm_metadata": {"guardrails": ["test_guardrail"]}, } - result = custom_guardrail.should_run_guardrail( - data=data, event_type=GuardrailEventHooks.pre_call - ) + result = custom_guardrail.should_run_guardrail(data=data, event_type=GuardrailEventHooks.pre_call) assert result is True @@ -240,9 +218,7 @@ class TestCustomGuardrailShouldRunGuardrail: "metadata": {"guardrails": ["test_guardrail"]}, } - result = custom_guardrail.should_run_guardrail( - data=data, event_type=GuardrailEventHooks.pre_call - ) + result = custom_guardrail.should_run_guardrail(data=data, event_type=GuardrailEventHooks.pre_call) assert result is True @@ -259,9 +235,7 @@ class TestCustomGuardrailShouldRunGuardrail: # Test with guardrails at root level data = {"model": "gpt-3.5-turbo", "guardrails": ["test_guardrail"]} - result = custom_guardrail.should_run_guardrail( - data=data, event_type=GuardrailEventHooks.pre_call - ) + result = custom_guardrail.should_run_guardrail(data=data, event_type=GuardrailEventHooks.pre_call) assert result is True @@ -281,9 +255,7 @@ class TestCustomGuardrailShouldRunGuardrail: "litellm_metadata": {"guardrails": ["different_guardrail"]}, } - result = custom_guardrail.should_run_guardrail( - data=data, event_type=GuardrailEventHooks.pre_call - ) + result = custom_guardrail.should_run_guardrail(data=data, event_type=GuardrailEventHooks.pre_call) assert result is False @@ -302,9 +274,7 @@ class TestCustomGuardrailShouldRunGuardrail: "model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": "test"}], } - result = custom_guardrail.should_run_guardrail( - data=data, event_type=GuardrailEventHooks.pre_call - ) + result = custom_guardrail.should_run_guardrail(data=data, event_type=GuardrailEventHooks.pre_call) assert result is True, "Global guardrail should run when default_on=True" # Test 2: User-injected disable at root level is IGNORED @@ -316,9 +286,7 @@ class TestCustomGuardrailShouldRunGuardrail: result = custom_guardrail.should_run_guardrail( data=data_with_disable_root, event_type=GuardrailEventHooks.pre_call ) - assert ( - result is True - ), "User-injected disable_global_guardrails should be ignored" + assert result is True, "User-injected disable_global_guardrails should be ignored" # Test 3: User-injected disable in metadata is IGNORED data_with_disable_metadata = { @@ -349,12 +317,8 @@ class TestCustomGuardrailShouldRunGuardrail: "metadata": {"user_api_key_metadata": {"disable_global_guardrails": True}}, "litellm_metadata": {"request_tags": ["user-supplied"]}, } - result = custom_guardrail.should_run_guardrail( - data=data_cross_key, event_type=GuardrailEventHooks.pre_call - ) - assert ( - result is False - ), "Admin config in metadata must not be shadowed by user-supplied litellm_metadata" + result = custom_guardrail.should_run_guardrail(data=data_cross_key, event_type=GuardrailEventHooks.pre_call) + assert result is False, "Admin config in metadata must not be shadowed by user-supplied litellm_metadata" # Test 6: After the pre-call strip runs, user-injected # user_api_key_metadata in the non-authoritative metadata key is gone. @@ -365,12 +329,8 @@ class TestCustomGuardrailShouldRunGuardrail: "metadata": {"user_api_key_metadata": {"disable_global_guardrails": True}}, "litellm_metadata": {}, # post-strip: attacker payload removed } - result = custom_guardrail.should_run_guardrail( - data=data_post_strip, event_type=GuardrailEventHooks.pre_call - ) - assert ( - result is False - ), "Admin config in metadata must be respected when other metadata key is empty" + result = custom_guardrail.should_run_guardrail(data=data_post_strip, event_type=GuardrailEventHooks.pre_call) + assert result is False, "Admin config in metadata must be respected when other metadata key is empty" def test_should_run_guardrail_key_disable_global_not_overruled_by_team_guardrail_list( self, @@ -436,12 +396,7 @@ class TestCustomGuardrailShouldRunGuardrail: "messages": [{"role": "user", "content": "test"}], "opted_out_global_guardrails": ["global_guardrail"], } - assert ( - custom_guardrail.should_run_guardrail( - data=data_root, event_type=GuardrailEventHooks.pre_call - ) - is True - ) + assert custom_guardrail.should_run_guardrail(data=data_root, event_type=GuardrailEventHooks.pre_call) is True # Test 2: User-injected opt-out in metadata is IGNORED data_metadata = { @@ -450,10 +405,7 @@ class TestCustomGuardrailShouldRunGuardrail: "metadata": {"opted_out_global_guardrails": ["global_guardrail"]}, } assert ( - custom_guardrail.should_run_guardrail( - data=data_metadata, event_type=GuardrailEventHooks.pre_call - ) - is True + custom_guardrail.should_run_guardrail(data=data_metadata, event_type=GuardrailEventHooks.pre_call) is True ) # Test 4: a different guardrail in the opt-out list → still runs @@ -462,12 +414,7 @@ class TestCustomGuardrailShouldRunGuardrail: "messages": [{"role": "user", "content": "test"}], "metadata": {"opted_out_global_guardrails": ["some_other_guardrail"]}, } - assert ( - custom_guardrail.should_run_guardrail( - data=data_other, event_type=GuardrailEventHooks.pre_call - ) - is True - ) + assert custom_guardrail.should_run_guardrail(data=data_other, event_type=GuardrailEventHooks.pre_call) is True # Test 5: empty opt-out list → still runs data_empty = { @@ -475,12 +422,7 @@ class TestCustomGuardrailShouldRunGuardrail: "messages": [{"role": "user", "content": "test"}], "metadata": {"opted_out_global_guardrails": []}, } - assert ( - custom_guardrail.should_run_guardrail( - data=data_empty, event_type=GuardrailEventHooks.pre_call - ) - is True - ) + assert custom_guardrail.should_run_guardrail(data=data_empty, event_type=GuardrailEventHooks.pre_call) is True # Test 6: malformed value (bool instead of list) → safely ignored, guardrail runs data_malformed = { @@ -489,10 +431,7 @@ class TestCustomGuardrailShouldRunGuardrail: "metadata": {"opted_out_global_guardrails": True}, } assert ( - custom_guardrail.should_run_guardrail( - data=data_malformed, event_type=GuardrailEventHooks.pre_call - ) - is True + custom_guardrail.should_run_guardrail(data=data_malformed, event_type=GuardrailEventHooks.pre_call) is True ) def test_should_run_guardrail_opt_out_does_not_affect_non_global(self): @@ -515,12 +454,69 @@ class TestCustomGuardrailShouldRunGuardrail: "guardrails": ["opt_in_guardrail"], }, } - assert ( - non_global.should_run_guardrail( - data=data, event_type=GuardrailEventHooks.pre_call - ) - is True + assert non_global.should_run_guardrail(data=data, event_type=GuardrailEventHooks.pre_call) 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.proxy.guardrails import auto_router_compression + from litellm.types.guardrails import GuardrailEventHooks + + always_on = CustomGuardrail( + guardrail_name="headroom-default", + default_on=True, + event_hook=GuardrailEventHooks.pre_call, ) + token = auto_router_compression._suppressed_compression_guardrails.set(frozenset({"headroom-default"})) + try: + assert ( + always_on.should_run_guardrail(data={"model": "smart-router"}, event_type=GuardrailEventHooks.pre_call) + is False + ) + finally: + auto_router_compression._suppressed_compression_guardrails.reset(token) + + def test_should_run_guardrail_suppression_does_not_affect_other_names(self): + from litellm.proxy.guardrails import auto_router_compression + from litellm.types.guardrails import GuardrailEventHooks + + always_on = CustomGuardrail( + guardrail_name="headroom-default", + default_on=True, + event_hook=GuardrailEventHooks.pre_call, + ) + token = auto_router_compression._suppressed_compression_guardrails.set(frozenset({"some-other-guardrail"})) + try: + assert ( + always_on.should_run_guardrail(data={"model": "smart-router"}, event_type=GuardrailEventHooks.pre_call) + is True + ) + finally: + auto_router_compression._suppressed_compression_guardrails.reset(token) + + def test_request_metadata_can_never_suppress_a_guardrail(self): + """Regression (security): suppression state is request-scoped and server-set, + never read from metadata. Metadata reaches spend logs the caller can read, so + anything honored from there is something a later request could replay to switch + off a PII or content-filter guardrail for itself.""" + from litellm.types.guardrails import GuardrailEventHooks + + always_on = CustomGuardrail( + guardrail_name="headroom-default", + default_on=True, + event_hook=GuardrailEventHooks.pre_call, + ) + forged = { + "model": "smart-router", + "metadata": { + "_auto_router_suppressed_compression_guardrails": [ + "headroom-default", + "any-token:headroom-default", + ], + }, + } + + assert always_on.should_run_guardrail(data=forged, event_type=GuardrailEventHooks.pre_call) is True class TestApplyGuardrailCheck: @@ -559,35 +555,33 @@ class TestApplyGuardrailCheck: child_with_override = ChildGuardrailWithOverride() # Test: CustomGuardrail itself has apply_guardrail in its __dict__ - assert ( - "apply_guardrail" in type(CustomGuardrail()).__dict__ - ), "CustomGuardrail should have apply_guardrail in its own __dict__" + assert "apply_guardrail" in type(CustomGuardrail()).__dict__, ( + "CustomGuardrail should have apply_guardrail in its own __dict__" + ) # Test: ParentGuardrail inherits but doesn't override, so it should NOT be in __dict__ - assert ( - "apply_guardrail" not in type(parent_instance).__dict__ - ), "ParentGuardrail should NOT have apply_guardrail in its own __dict__ (only inherited)" + assert "apply_guardrail" not in type(parent_instance).__dict__, ( + "ParentGuardrail should NOT have apply_guardrail in its own __dict__ (only inherited)" + ) # Test: ChildGuardrailWithoutOverride only inherits, should NOT be in __dict__ - assert ( - "apply_guardrail" not in type(child_without_override).__dict__ - ), "ChildGuardrailWithoutOverride should NOT have apply_guardrail in its own __dict__ (only inherited)" + assert "apply_guardrail" not in type(child_without_override).__dict__, ( + "ChildGuardrailWithoutOverride should NOT have apply_guardrail in its own __dict__ (only inherited)" + ) # Test: ChildGuardrailWithOverride overrides the method, SHOULD be in __dict__ - assert ( - "apply_guardrail" in type(child_with_override).__dict__ - ), "ChildGuardrailWithOverride SHOULD have apply_guardrail in its own __dict__ (overridden)" + assert "apply_guardrail" in type(child_with_override).__dict__, ( + "ChildGuardrailWithOverride SHOULD have apply_guardrail in its own __dict__ (overridden)" + ) # Verify that all instances still have the method via inheritance (hasattr) - assert hasattr( - parent_instance, "apply_guardrail" - ), "All instances should have apply_guardrail via inheritance" - assert hasattr( - child_without_override, "apply_guardrail" - ), "All instances should have apply_guardrail via inheritance" - assert hasattr( - child_with_override, "apply_guardrail" - ), "All instances should have apply_guardrail via inheritance" + assert hasattr(parent_instance, "apply_guardrail"), "All instances should have apply_guardrail via inheritance" + assert hasattr(child_without_override, "apply_guardrail"), ( + "All instances should have apply_guardrail via inheritance" + ) + assert hasattr(child_with_override, "apply_guardrail"), ( + "All instances should have apply_guardrail via inheritance" + ) class TestGuardrailLoggingAggregation: @@ -614,11 +608,7 @@ class TestGuardrailLoggingAggregation: def test_appends_to_existing_metadata_list(self): request_data = { - "metadata": { - "standard_logging_guardrail_information": [ - {"guardrail_name": "existing_guardrail"} - ] - } + "metadata": {"standard_logging_guardrail_information": [{"guardrail_name": "existing_guardrail"}]} } self._invoke_add_log(request_data) @@ -630,11 +620,7 @@ class TestGuardrailLoggingAggregation: assert info[1]["guardrail_name"] == "test_guardrail" def test_converts_existing_metadata_dict_to_list(self): - request_data = { - "metadata": { - "standard_logging_guardrail_information": {"guardrail_name": "legacy"} - } - } + request_data = {"metadata": {"standard_logging_guardrail_information": {"guardrail_name": "legacy"}}} self._invoke_add_log(request_data) @@ -646,18 +632,12 @@ class TestGuardrailLoggingAggregation: def test_appends_to_litellm_metadata(self): request_data = { - "litellm_metadata": { - "standard_logging_guardrail_information": [ - {"guardrail_name": "litellm_existing"} - ] - } + "litellm_metadata": {"standard_logging_guardrail_information": [{"guardrail_name": "litellm_existing"}]} } self._invoke_add_log(request_data) - info = request_data["litellm_metadata"][ - "standard_logging_guardrail_information" - ] + info = request_data["litellm_metadata"]["standard_logging_guardrail_information"] assert isinstance(info, list) assert len(info) == 2 assert info[1]["guardrail_name"] == "test_guardrail" @@ -674,12 +654,10 @@ class TestGuardrailLoggingAggregation: self._invoke_add_log(request_data) - assert ( - "standard_logging_guardrail_information" not in request_data["metadata"] - ), "entry landed in the caller's metadata, where the spend log does not read it" - info = request_data["litellm_metadata"][ - "standard_logging_guardrail_information" - ] + assert "standard_logging_guardrail_information" not in request_data["metadata"], ( + "entry landed in the caller's metadata, where the spend log does not read it" + ) + info = request_data["litellm_metadata"]["standard_logging_guardrail_information"] assert len(info) == 1 assert info[0]["guardrail_name"] == "test_guardrail" @@ -697,9 +675,7 @@ class TestGuardrailLoggingAggregation: } self._invoke_add_log(request_data) - add_guardrail_to_applied_guardrails_header( - request_data=request_data, guardrail_name="test_guardrail" - ) + add_guardrail_to_applied_guardrails_header(request_data=request_data, guardrail_name="test_guardrail") buckets = { key @@ -745,9 +721,7 @@ class TestGuardrailOtelSpanEmission: assert len(captured) == 1 emitted = captured[0] - recorded = request_data["metadata"]["standard_logging_guardrail_information"][ - -1 - ] + recorded = request_data["metadata"]["standard_logging_guardrail_information"][-1] assert emitted is recorded assert emitted["guardrail_name"] == "emit_guard" assert emitted["start_time"] == 1.0 @@ -757,9 +731,7 @@ class TestGuardrailOtelSpanEmission: def _boom(_entry): raise RuntimeError("otel exporter down") - monkeypatch.setattr( - "litellm.integrations.otel.logger.emit_guardrail_span", _boom - ) + monkeypatch.setattr("litellm.integrations.otel.logger.emit_guardrail_span", _boom) request_data = {"metadata": {}} self._record(self._make_guardrail(), request_data) @@ -856,9 +828,7 @@ class TestGuardrailSensitiveFieldStripping: duration=1.0, ) - logged_response = request_data["metadata"][ - "standard_logging_guardrail_information" - ][0]["guardrail_response"] + logged_response = request_data["metadata"]["standard_logging_guardrail_information"][0]["guardrail_response"] assert "secret_fields" not in logged_response assert "sk-live-SHOULD-NOT-APPEAR" not in json.dumps(logged_response) @@ -871,9 +841,7 @@ class TestGuardrailSensitiveFieldStripping: guardrail_json_response=[ { "result": "ok", - "secret_fields": { - "raw_headers": {"authorization": "Bearer sk-secret"} - }, + "secret_fields": {"raw_headers": {"authorization": "Bearer sk-secret"}}, }, {"result": "also_ok"}, ], @@ -927,9 +895,7 @@ class TestGuardrailResponseCredentialMasking: duration=1.0, ) - logged = request_data["metadata"]["standard_logging_guardrail_information"][0][ - "guardrail_response" - ] + logged = request_data["metadata"]["standard_logging_guardrail_information"][0]["guardrail_response"] masked_key = logged["metadata_snapshot"]["callback_vars"]["langsmith_api_key"] assert masked_key != plaintext_key @@ -938,10 +904,7 @@ class TestGuardrailResponseCredentialMasking: assert logged["model"] == "gpt-4o-mini" assert logged["messages"] == [{"role": "user", "content": "hi"}] - assert ( - logged["metadata_snapshot"]["callback_vars"]["langsmith_project"] - == "proj-name" - ) + assert logged["metadata_snapshot"]["callback_vars"]["langsmith_project"] == "proj-name" def test_nested_user_api_key_auth_metadata_is_masked(self): import json @@ -1000,9 +963,7 @@ class TestGuardrailResponseCredentialMasking: request_data: dict = {"metadata": {}} guardrail.add_standard_logging_guardrail_information_to_request_data( - guardrail_json_response={ - "filters": [{"regex": r"\d{3}-\d{2}-\d{4}", "action": "BLOCKED"}] - }, + guardrail_json_response={"filters": [{"regex": r"\d{3}-\d{2}-\d{4}", "action": "BLOCKED"}]}, request_data=request_data, guardrail_status="success", ) @@ -1025,9 +986,7 @@ class TestGuardrailResponseCredentialMasking: guardrail_status="success", ) - logged = request_data["metadata"]["standard_logging_guardrail_information"][0][ - "guardrail_response" - ] + logged = request_data["metadata"]["standard_logging_guardrail_information"][0]["guardrail_response"] assert logged["flagged"] is True assert logged["score"] == 0.94 assert logged["tokens_used"] == 42 @@ -1039,18 +998,14 @@ class TestGuardrailResponseCredentialMasking: plaintext = "lsv2_pt_abcdef1234567890" guardrail.add_standard_logging_guardrail_information_to_request_data( - guardrail_json_response={ - "metadata_snapshot": { - "callback_vars": {"langsmith_api_key": plaintext} - } - }, + guardrail_json_response={"metadata_snapshot": {"callback_vars": {"langsmith_api_key": plaintext}}}, request_data=request_data, guardrail_status="success", ) - masked = request_data["metadata"]["standard_logging_guardrail_information"][0][ - "guardrail_response" - ]["metadata_snapshot"]["callback_vars"]["langsmith_api_key"] + masked = request_data["metadata"]["standard_logging_guardrail_information"][0]["guardrail_response"][ + "metadata_snapshot" + ]["callback_vars"]["langsmith_api_key"] assert masked != plaintext assert masked.startswith(plaintext[:4]) assert masked.endswith(plaintext[-4:]) @@ -1544,9 +1499,7 @@ class TestEventTypeLogging: guardrail = TestGuardrail() request_data = {"metadata": {}} - await guardrail.apply_guardrail( - inputs={"texts": ["x"]}, request_data=request_data - ) + await guardrail.apply_guardrail(inputs={"texts": ["x"]}, request_data=request_data) logged_info = request_data["metadata"]["standard_logging_guardrail_information"] assert len(logged_info) == 1, ( @@ -1588,9 +1541,7 @@ class TestEventTypeLogging: request_data = {"metadata": {}} with pytest.raises(ValueError, match="blocked"): - await guardrail.apply_guardrail( - inputs={"texts": ["x"]}, request_data=request_data - ) + await guardrail.apply_guardrail(inputs={"texts": ["x"]}, request_data=request_data) logged_info = request_data["metadata"]["standard_logging_guardrail_information"] assert len(logged_info) == 1 @@ -1719,9 +1670,7 @@ class TestTracingFieldsPopulation: guardrail_json_response="blocked", request_data=request_data, guardrail_status="guardrail_intervened", - tracing_detail=GuardrailTracingDetail( - policy_template="EU AI Act Article 5" - ), + tracing_detail=GuardrailTracingDetail(policy_template="EU AI Act Article 5"), ) slg_list = request_data["metadata"]["standard_logging_guardrail_information"] @@ -1763,13 +1712,7 @@ class TestCustomGuardrailSpendLogMatchRedaction: cg = CustomGuardrail(guardrail_name="test-rail") raw = { "assessments": [ - { - "sensitiveInformationPolicy": { - "piiEntities": [ - {"type": "NAME", "match": "GG", "action": "BLOCKED"} - ] - } - } + {"sensitiveInformationPolicy": {"piiEntities": [{"type": "NAME", "match": "GG", "action": "BLOCKED"}]}} ] } request_data: dict = {"metadata": {}} @@ -1780,17 +1723,10 @@ class TestCustomGuardrailSpendLogMatchRedaction: ) slg = request_data["metadata"]["standard_logging_guardrail_information"][0] assert ( - slg["guardrail_response"]["assessments"][0]["sensitiveInformationPolicy"][ - "piiEntities" - ][0]["match"] + slg["guardrail_response"]["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"][0]["match"] == "[REDACTED]" ) - assert ( - raw["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"][0][ - "match" - ] - == "GG" - ) + assert raw["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"][0]["match"] == "GG" def test_add_standard_logging_redacts_regex_field(self): cg = CustomGuardrail(guardrail_name="test-rail") 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..db2f94306fb --- /dev/null +++ b/tests/test_litellm/proxy/guardrails/test_auto_router_compression.py @@ -0,0 +1,381 @@ +"""Unit tests for litellm.proxy.guardrails.auto_router_compression.""" + +import json +from typing import Any + +import pytest + +from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.proxy.guardrails import auto_router_compression +from litellm.proxy.guardrails.auto_router_compression import ( + AutoRouterCompressionPolicy, + arm_pre_call, + messages_for_routing, + policy_for_model, + policy_from_litellm_params, +) +from litellm.types.guardrails import GuardrailEventHooks +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] + + +def _marker(compression: dict[str, str], tags: list[str] | None = None) -> dict[str, Any]: + return { + "model_name": "smart-router", + "litellm_params": { + "model": "auto_router/complexity_router", + **compression, + **({"tags": tags} if tags is not None else {}), + }, + } + + +class TestPolicyForModel: + def test_no_router_returns_none(self): + assert policy_for_model(llm_router=None, model_alias="smart-router", team_id=None, request_tags=()) 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, request_tags=()) 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, request_tags=()) is None + + def test_marker_deployment_with_policy_is_found(self): + router = _FakeRouter( + [_marker({"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, request_tags=()) + assert policy == AutoRouterCompressionPolicy(routing="headroom-a", model=None) + + def test_picks_the_marker_whose_tags_the_request_carries(self): + """Regression: an alias with several tag-scoped markers must not suppress one + marker's guardrail and then route under a different marker's policy.""" + router = _FakeRouter( + [ + _marker({"auto_router_routing_compression": "headroom-eu"}, tags=["eu"]), + _marker({"auto_router_routing_compression": "headroom-us"}, tags=["us"]), + ] + ) + + eu = policy_for_model(llm_router=router, model_alias="smart-router", team_id=None, request_tags=("eu",)) + us = policy_for_model(llm_router=router, model_alias="smart-router", team_id=None, request_tags=("us",)) + + assert eu == AutoRouterCompressionPolicy(routing="headroom-eu", model=None) + assert us == AutoRouterCompressionPolicy(routing="headroom-us", model=None) + + def test_untagged_marker_matches_any_request(self): + router = _FakeRouter([_marker({"auto_router_routing_compression": "headroom-a"})]) + policy = policy_for_model( + llm_router=router, model_alias="smart-router", team_id=None, request_tags=("anything",) + ) + assert policy == AutoRouterCompressionPolicy(routing="headroom-a", model=None) + + def test_a_marker_scoped_to_other_tags_is_never_the_fallback(self): + """Regression: a "us" request must not fall back to an "eu" marker's policy.""" + router = _FakeRouter( + [ + _marker({"auto_router_routing_compression": "headroom-eu"}, tags=["eu"]), + _marker({"auto_router_routing_compression": "headroom-default"}), + ] + ) + policy = policy_for_model(llm_router=router, model_alias="smart-router", team_id=None, request_tags=("us",)) + assert policy == AutoRouterCompressionPolicy(routing="headroom-default", model=None) + + def test_no_untagged_fallback_means_no_policy(self): + """No matching marker means no policy, not an unrelated slice's compression.""" + router = _FakeRouter([_marker({"auto_router_routing_compression": "headroom-eu"}, tags=["eu"])]) + assert ( + policy_for_model(llm_router=router, model_alias="smart-router", team_id=None, request_tags=("us",)) is None + ) + + def test_tag_scoped_marker_takes_precedence_over_untagged(self): + """Regression: when multiple markers exist, the tag-scoped one the request + actually matches should be used, not the first untagged one.""" + router = _FakeRouter( + [ + _marker({"auto_router_routing_compression": "headroom-untagged"}), + _marker({"auto_router_routing_compression": "headroom-eu"}, tags=["eu"]), + ] + ) + policy = policy_for_model(llm_router=router, model_alias="smart-router", team_id=None, request_tags=("eu",)) + assert policy == AutoRouterCompressionPolicy(routing="headroom-eu", 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(monkeypatch): + import litellm + from litellm.proxy.guardrails import guardrail_registry + + # Registered under a compression provider name: both hops refuse a name that does + # not resolve to one, so a bare callback would (correctly) never be used. + monkeypatch.setitem(guardrail_registry.guardrail_class_registry, "headroom", _RecordingCompressionGuardrail) + 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 _NonCompressionGuardrail(CustomGuardrail): + """A guardrail that is not a compression provider, e.g. a PII or content filter.""" + + def __init__(self, guardrail_name: str): + super().__init__(guardrail_name=guardrail_name) + self.called = False + + async def apply_guardrail( + self, inputs: GenericGuardrailAPIInputs, request_data: dict, input_type: str, logging_obj=None + ) -> GenericGuardrailAPIInputs: + self.called = True + return inputs + + +class TestArmPreCall: + @pytest.mark.asyncio + async def test_no_router_is_noop(self): + data = {"model": "smart-router", "messages": [{"role": "user", "content": "hi"}]} + await arm_pre_call(data=data, llm_router=None) + assert "metadata" not in data + + @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"}]} + await arm_pre_call(data=data, llm_router=router) + assert "metadata" not in data + assert "litellm_metadata" not in data + + @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"}]} + await arm_pre_call(data=data, llm_router=router) + assert auto_router_compression.suppressed_compression_guardrails() == frozenset({"always-on-compression"}) + # Suppression state must never ride along in metadata: that reaches spend + # logs the caller can read, and anything there is replayable. + assert "always-on-compression" not in json.dumps(data.get("metadata", {})) + assert always_on.should_run_guardrail(data=data, event_type=GuardrailEventHooks.pre_call) is False + finally: + litellm.logging_callback_manager.remove_callback_from_all_lists(always_on) + + @pytest.mark.asyncio + async def test_suppression_state_never_enters_request_metadata(self): + """Regression (security): metadata reaches spend logs, so a suppression list + there is one a caller could read back and replay to disable a guardrail.""" + guardrail = _RecordingCompressionGuardrail(guardrail_name="always-on-compression") + import litellm + + litellm.logging_callback_manager.add_litellm_callback(guardrail) + try: + router = _FakeRouter([_marker({"auto_router_routing_compression": "headroom-a"})]) + data = {"model": "smart-router", "messages": [{"role": "user", "content": "hi"}]} + await arm_pre_call(data=data, llm_router=router) + assert "suppress" not in json.dumps(data).lower() + finally: + litellm.logging_callback_manager.remove_callback_from_all_lists(guardrail) + + @pytest.mark.asyncio + async def test_model_side_guardrail_is_requested_even_when_not_default_on(self, monkeypatch): + import litellm + from litellm.proxy.guardrails import guardrail_registry + + monkeypatch.setitem(guardrail_registry.guardrail_class_registry, "headroom", _RecordingCompressionGuardrail) + active = _RecordingCompressionGuardrail(guardrail_name="headroom-b") + litellm.logging_callback_manager.add_litellm_callback(active) + 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"}]} + try: + await arm_pre_call(data=data, llm_router=router) + assert data["metadata"]["guardrails"] == ["headroom-b"] + finally: + litellm.logging_callback_manager.remove_callback_from_all_lists(active) + + @pytest.mark.asyncio + async def test_arm_pre_call_keeps_no_copy_of_the_prompt(self): + """Regression (security): arm_pre_call runs before the guardrails, so any copy it + kept would be pre-masking text that routing then POSTs to an external service.""" + router = _FakeRouter([_marker({"auto_router_routing_compression": "headroom-a"})]) + data = {"model": "smart-router", "messages": [{"role": "user", "content": "my ssn is 123-45-6789"}]} + + await arm_pre_call(data=data, llm_router=router) + + assert "123-45-6789" not in json.dumps(data.get("metadata", {})) + assert not hasattr(auto_router_compression, "_routing_messages_snapshot") + + +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_none_with_no_model_compression_returns_none(self): + """Nothing compressed either hop, so the caller's own messages are already right.""" + policy = AutoRouterCompressionPolicy(routing=None, model=None) + assert await messages_for_routing(policy=policy, messages=[], request_kwargs={}) is None + + @pytest.mark.asyncio + async def test_routing_none_never_reaches_for_a_pre_guardrail_copy(self): + """No uncompressed copy survives the model hop, and keeping one would mean + retaining the pre-masking text. Routing reads what it has.""" + policy = AutoRouterCompressionPolicy(routing=None, model="headroom-a") + model_compressed = [{"role": "user", "content": "[COMPRESSED] the full original conversation"}] + + assert await messages_for_routing(policy=policy, messages=model_compressed, request_kwargs={}) is None + + @pytest.mark.asyncio + async def test_unknown_guardrail_name_routes_on_the_uncompressed_messages(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 == messages + + @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_routing_compresses_what_the_other_guardrails_left_behind(self, registered_guardrail): + """Regression (security): routing POSTs its input out, so it must read what the + earlier guardrails left behind, not a pre-masking copy.""" + policy = AutoRouterCompressionPolicy(routing="fake-compress", model="headroom-b") + masked = [{"role": "user", "content": "my ssn is [REDACTED]"}] + + result = await messages_for_routing(policy=policy, messages=masked, request_kwargs={}) + + assert result == [{"role": "user", "content": "[COMPRESSED] my ssn is [REDACTED]"}] + assert registered_guardrail.request_data_seen[0]["messages"] == masked + + @pytest.mark.asyncio + async def test_a_non_compression_guardrail_is_never_invoked_for_routing(self, monkeypatch): + """Regression (security): naming an ordinary guardrail must not turn the routing + hop into a way to ship prompts to whatever service backs it.""" + import litellm + + other = _NonCompressionGuardrail(guardrail_name="pii-filter") + litellm.logging_callback_manager.add_litellm_callback(other) + try: + policy = AutoRouterCompressionPolicy(routing="pii-filter", model=None) + messages = [{"role": "user", "content": "my ssn is 123-45-6789"}] + + result = await messages_for_routing(policy=policy, messages=messages, request_kwargs={}) + + assert other.called is False + assert result == messages + finally: + litellm.logging_callback_manager.remove_callback_from_all_lists(other) + + @pytest.mark.asyncio + async def test_guardrail_receives_a_throwaway_request_data_not_the_real_request_kwargs(self, registered_guardrail): + """Regression: a guardrail writes stats onto the request_data it is given, so + passing the caller's own would double-count into extract_compression_saved_tokens.""" + 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..96d9b0c5a26 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -376,6 +376,75 @@ 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.""" + from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.proxy.guardrails import guardrail_registry + + # The model hop is only armed for a name that resolves to an active compression + # guardrail, so arming it has to have a real one to resolve to. + class _FakeCompressionGuardrail(CustomGuardrail): + pass + + monkeypatch.setitem(guardrail_registry.guardrail_class_registry, "headroom", _FakeCompressionGuardrail) + active_guardrail = _FakeCompressionGuardrail(guardrail_name="headroom-model") + litellm.logging_callback_manager.add_litellm_callback(active_guardrail) + + 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) + + try: + 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, + ) + finally: + litellm.logging_callback_manager.remove_callback_from_all_lists(active_guardrail) + + 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 8368aa11316..7c044310e14 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -20,6 +20,7 @@ import respx 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.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.llms.bedrock.common_utils import BedrockError @@ -9996,6 +9997,215 @@ 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, monkeypatch): + from litellm.proxy.guardrails import guardrail_registry + + # Registered under a compression provider name: both hops refuse a name that + # does not resolve to one, so a bare callback would never be used. + monkeypatch.setitem(guardrail_registry.guardrail_class_registry, "headroom", self._CompressingGuardrail) + 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_routing_none_classifies_on_the_live_messages_not_a_pre_guardrail_copy(self, registered_guardrail): + """Routing asked for no compression while the model hop compressed, so the only + messages left are that guardrail's output and the strategy classifies on them. + + Keeping a pre-compression copy to classify on instead is what this deliberately + gives up: that copy is taken before the pre-call guardrails run, so it still + holds whatever a masking guardrail exists to strip, and routing-side compression + POSTs its input to an external service.""" + router, strategy = self._router( + { + "auto_router_routing_compression": "none", + "auto_router_model_compression": "fake-compress", + } + ) + model_compressed = [{"role": "user", "content": "[COMPRESSED] What is the capital of France?"}] + + await router.async_pre_routing_hook( + model="smart-router", request_kwargs={"metadata": {}}, messages=model_compressed + ) + + assert strategy.received_messages == model_compressed + assert registered_guardrail.call_count == 0 + + @pytest.mark.asyncio + @pytest.mark.asyncio + async def test_same_compression_still_compresses_routing_when_nothing_armed_it(self, registered_guardrail): + """Regression: only the proxy calls arm_pre_call. Used through the SDK, nothing + arms the model-side guardrail and nothing has compressed anything, so reusing a + model-hop result that was never produced would serve the request with no + compression on either hop, silently ignoring the configuration.""" + from litellm.proxy.guardrails import auto_router_compression + + router, strategy = self._router( + { + "auto_router_routing_compression": "fake-compress", + "auto_router_model_compression": "fake-compress", + } + ) + uncompressed = self._messages() + assert auto_router_compression.model_hop_compression_armed() is False + + await router.async_pre_routing_hook( + model="smart-router", request_kwargs={"metadata": {}}, messages=uncompressed + ) + + assert strategy.received_messages != uncompressed + assert registered_guardrail.call_count == 1 + + 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.""" + from litellm.proxy.guardrails import auto_router_compression + + 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?"}] + # arm_pre_call is what would have armed that guardrail, and only the proxy calls + # it; the reuse below is conditional on it having run. + armed = auto_router_compression._model_hop_armed.set(True) + + try: + response = await router.async_pre_routing_hook( + model="smart-router", request_kwargs={"metadata": {}}, messages=already_compressed_messages + ) + finally: + auto_router_compression._model_hop_armed.reset(armed) + + 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") + @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 cad41557100..d7df6ce33bb 100644 --- a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx @@ -1,11 +1,11 @@ import { SimpleTooltip } from "@/components/ui/tooltip"; import { MultiSelect } from "@/components/shared/MultiSelect"; import { SearchSelect } from "@/components/shared/SearchSelect"; -import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; import { ChevronRight, Info, Plus, Trash2, X } from "lucide-react"; import { Switch } from "@/components/ui/switch"; import { AffinityControls } from "./AffinityControls"; +import TierRowSelect from "./TierRowSelect"; import { ModalityRoutingControls } from "./ModalityRoutingControls"; import { Card, CardContent } from "@/components/ui/card"; import { Collapsible, CollapsibleContent, CollapsibleTrigger } from "@/components/ui/collapsible"; @@ -50,6 +50,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"; @@ -370,27 +372,6 @@ const TierRowEditFields: React.FC<{ ); -const TierRowSelect: React.FC<{ - label: string; - options: { value: string; label: string }[]; - value: string | null; - onValueChange: (rowId: string) => void; - placeholder?: string; -}> = ({ label, options, value, onValueChange, placeholder }) => ( - -); - export type AdaptiveEligible = "all" | "classified_tier"; export type ComplexityTierLabels = Partial>; @@ -502,6 +483,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; } @@ -604,6 +589,8 @@ const ComplexityRouterConfig: React.FC = ({ onMatchThresholdChange = () => {}, escalationKeywords = [], onEscalationKeywordsChange, + autoRouterCompression = DEFAULT_AUTO_ROUTER_COMPRESSION, + onAutoRouterCompressionChange, showValidationErrors = false, }) => { const customTierSet = value.custom_tier_set; @@ -877,6 +864,17 @@ const ComplexityRouterConfig: React.FC = ({ }, ] : []), + ...(onAutoRouterCompressionChange + ? [ + { + key: "compression", + label: Advanced: Compression, + children: ( + + ), + }, + ] + : []), ...(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..c1817918f60 --- /dev/null +++ b/ui/litellm-dashboard/src/components/add_model/CompressionControls.tsx @@ -0,0 +1,91 @@ +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 { + AutoRouterCompressionState, + isCompressionGuardrailProvider, + NO_COMPRESSION, +} from "./buildAutoRouterCompression"; + +interface CompressionControlsProps { + value: AutoRouterCompressionState; + onChange: (state: AutoRouterCompressionState) => void; +} + +const NONE_OPTION: SearchSelectOption = { label: "None (no compression)", value: NO_COMPRESSION }; + +const CompressionControls: React.FC = ({ value, onChange }) => { + const { routing, sameAsRouting, model } = value; + const onRoutingChange = (newRouting: string | undefined) => + onChange({ ...value, routing: newRouting, sameAsRouting: newRouting === undefined ? true : sameAsRouting }); + const onSameAsRoutingChange = (newSameAsRouting: boolean) => onChange({ ...value, sameAsRouting: newSameAsRouting }); + const onModelChange = (newModel: string | undefined) => onChange({ ...value, model: newModel }); + + const { data } = useGuardrails(); + const compressionOptions: SearchSelectOption[] = (data?.guardrails ?? []) + .filter((g) => isCompressionGuardrailProvider(g.litellm_params?.guardrail)) + .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/TierRowSelect.tsx b/ui/litellm-dashboard/src/components/add_model/TierRowSelect.tsx new file mode 100644 index 00000000000..ad7d53f9eae --- /dev/null +++ b/ui/litellm-dashboard/src/components/add_model/TierRowSelect.tsx @@ -0,0 +1,25 @@ +import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; +import React from "react"; + +const TierRowSelect: React.FC<{ + label: string; + options: { value: string; label: string }[]; + value: string | null; + onValueChange: (rowId: string) => void; + placeholder?: string; +}> = ({ label, options, value, onValueChange, placeholder }) => ( + +); + +export default TierRowSelect; 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 ff71f161800..60780f9cbe0 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); @@ -465,6 +473,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); @@ -670,6 +679,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..20d6af50d18 --- /dev/null +++ b/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.test.ts @@ -0,0 +1,119 @@ +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 no model-hop compression, not same-as-routing", () => { + const state = hydrateAutoRouterCompression({ auto_router_routing_compression: "headroom-a" }); + expect(state).toEqual({ routing: "headroom-a", sameAsRouting: false, model: "none" }); + }); + + it("re-saving a routing-only config leaves the model hop uncompressed", () => { + // Regression: the backend reads an absent model key as no model-hop compression. + // Hydrating it as same-as-routing made opening the router and saving any unrelated + // edit write the routing guardrail onto the model hop, so the model call silently + // started receiving compressed messages. + const stored = { auto_router_routing_compression: "headroom-a" }; + const rebuilt = buildAutoRouterCompressionParams(hydrateAutoRouterCompression(stored)); + expect(rebuilt.auto_router_model_compression).toBe("none"); + expect(rebuilt.auto_router_model_compression).not.toBe("headroom-a"); + }); + + it("surfaces a stored model-only policy instead of reading as untouched", () => { + // Regression: the backend treats either key alone as an authoritative policy, so a + // model-only config that hydrated to the inherit state was invisible in the form, + // and the next save overwrote the stored model hop with the routing value. + const state = hydrateAutoRouterCompression({ auto_router_model_compression: "headroom-b" }); + expect(state).toEqual({ routing: "none", sameAsRouting: false, model: "headroom-b" }); + }); + + it("round-trips a model-only policy without changing either hop", () => { + const stored = { auto_router_model_compression: "headroom-b" }; + const rebuilt = buildAutoRouterCompressionParams(hydrateAutoRouterCompression(stored)); + expect(rebuilt.auto_router_model_compression).toBe("headroom-b"); + expect(rebuilt.auto_router_routing_compression).toBe("none"); + }); + + 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..6f401a12865 --- /dev/null +++ b/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.ts @@ -0,0 +1,69 @@ +/** + * 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"; + +/** Guardrail providers that compress prompts, mirroring COMPRESSION_GUARDRAIL_PROVIDERS in + * litellm/proxy/guardrails/auto_router_compression.py. Both are selectable per hop. */ +export const COMPRESSION_GUARDRAIL_PROVIDERS: readonly string[] = ["headroom", "compresr"]; + +export const isCompressionGuardrailProvider = (provider: unknown): boolean => + typeof provider === "string" && COMPRESSION_GUARDRAIL_PROVIDERS.includes(provider.toLowerCase()); + +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 storedRouting = litellmParams.auto_router_routing_compression ?? undefined; + const storedModel = litellmParams.auto_router_model_compression ?? undefined; + + // Only neither key set means the section was never touched. The backend treats + // either key on its own as an authoritative policy (policy_from_litellm_params), so + // reading a model-only config as untouched would hide it from the form and let the + // next save overwrite the stored model hop. + if (storedRouting === undefined && storedModel === undefined) return DEFAULT_AUTO_ROUTER_COMPRESSION; + + // An absent key on either hop is no compression for that hop, not same-as-the-other: + // the backend reads it as None. Hydrating it as same-as-routing would make re-saving + // an unrelated edit write one hop's guardrail onto the other. + const routing = storedRouting ?? NO_COMPRESSION; + const model = storedModel ?? NO_COMPRESSION; + const sameAsRouting = 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 e45f84fa646..3367c810991 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 @@ -1036,6 +1036,87 @@ 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"); + }); +}); describe("EditAutoRouterModal classifier vision", () => { beforeEach(() => { 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 dd3ff1241ea..3c0013e267e 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, @@ -447,6 +453,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", @@ -539,6 +548,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, @@ -651,6 +666,7 @@ const EditAutoRouterModal: React.FC = ({ ...modelData.litellm_params, complexity_router_config: updatedConfig, complexity_router_default_model: defaultModel, + ...buildAutoRouterCompressionParams(autoRouterCompression), }; const updatedModelInfo = { ...modelData.model_info, @@ -772,6 +788,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 14485a0a115..c81ea2d7874 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -29281,6 +29281,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 */ @@ -39411,6 +39415,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 */