From dd60b7e40f44d7687ad4577332d6f57fc21d7af3 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Fri, 4 Sep 2026 16:16:38 -0700 Subject: [PATCH 01/13] feat(auto-router): decouple compression between the routing decision and the model call An auto router marker deployment can now set auto_router_routing_compression and auto_router_model_compression in its litellm_params, naming the compression guardrail each hop should use (or "none" for no compression on that hop). Neither key set means the request's own compression guardrails keep applying to both hops unchanged. Backend: Router.async_pre_routing_hook resolves the marker's policy and compresses a copy of the messages for the routing decision only when the policy differs from what the model call already got; when both hops share the same compression, it reuses what the ordinary pre-call guardrail pipeline already produced instead of compressing twice. The proxy layer suppresses every other compression guardrail once a policy is engaged and arms the model-side guardrail even when it is not default_on. UI: the auto router's Detailed Configuration gains an Advanced: Compression section with a routing-decision selector and a same/different toggle for the model call, matching the same/different address pattern. --- litellm/constants.py | 4 + litellm/integrations/custom_guardrail.py | 18 ++ litellm/proxy/common_request_processing.py | 7 + .../guardrails/auto_router_compression.py | 211 +++++++++++++ litellm/router.py | 53 +++- litellm/types/router.py | 4 + litellm/types/utils.py | 2 + .../integrations/test_custom_guardrail.py | 44 +++ .../test_auto_router_compression.py | 285 ++++++++++++++++++ .../proxy/test_common_request_processing.py | 54 ++++ tests/test_litellm/test_router.py | 151 ++++++++++ .../add_model/ComplexityRouterConfig.tsx | 34 +++ .../add_model/CompressionControls.tsx | 93 ++++++ .../add_model/add_auto_router_tab.test.tsx | 75 ++++- .../add_model/add_auto_router_tab.tsx | 11 + .../buildAutoRouterCompression.test.ts | 93 ++++++ .../add_model/buildAutoRouterCompression.ts | 52 ++++ .../handle_add_auto_router_submit.tsx | 5 +- .../edit_auto_router_modal.test.tsx | 81 +++++ .../edit_auto_router_modal.tsx | 18 ++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 8 + 21 files changed, 1297 insertions(+), 6 deletions(-) create mode 100644 litellm/proxy/guardrails/auto_router_compression.py create mode 100644 tests/test_litellm/proxy/guardrails/test_auto_router_compression.py create mode 100644 ui/litellm-dashboard/src/components/add_model/CompressionControls.tsx create mode 100644 ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.test.ts create mode 100644 ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.ts diff --git a/litellm/constants.py b/litellm/constants.py index da731cb5eb2..25fdaec20de 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -215,6 +215,10 @@ MAX_CALLBACKS: Final = get_env_int("LITELLM_MAX_CALLBACKS", 100) # so the deployment-level hook does not re-run them for the same request PRE_CALL_EXECUTED_GUARDRAILS_KEY: Final = "_pre_call_executed_guardrails" +# Metadata key listing compression guardrails an auto router's own compression +# policy suppresses for this request. See litellm.proxy.guardrails.auto_router_compression. +AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY: Final = "_auto_router_suppressed_compression_guardrails" + # Generic fallback for unknown models DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET: Final = int( os.getenv("DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET", 128) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 9bb613654e4..7f8effa2317 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -45,6 +45,7 @@ dc: Final = DualCache() from litellm.constants import ( + AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY, GUARDRAIL_SCANNED_MESSAGES_CACHE_TTL_SECONDS, PRE_CALL_EXECUTED_GUARDRAILS_KEY, ) @@ -940,6 +941,20 @@ class CustomGuardrail(CustomLogger): """ return False + def _suppressed_by_auto_router_compression(self, data: dict) -> bool: + """True when an auto router's own compression policy suppresses this guardrail. + + Set only by litellm.proxy.guardrails.auto_router_compression.arm_pre_call, never + by the caller, so a request cannot suppress its own guardrails this way. + """ + for meta_key in ("metadata", "litellm_metadata"): + meta = data.get(meta_key) + if isinstance(meta, dict): + suppressed = meta.get(AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY) + if isinstance(suppressed, list) and self.guardrail_name in suppressed: + return True + return False + def should_run_guardrail( self, data, @@ -948,6 +963,9 @@ class CustomGuardrail(CustomLogger): """ Returns True if the guardrail should be run on the event_type """ + if self._suppressed_by_auto_router_compression(data): + return False + requested_guardrails: Final = self.get_guardrail_from_metadata(data) disable_global_guardrail: Final = self.get_disable_global_guardrail(data) opted_out_global_guardrails: Final = self.get_opted_out_global_guardrails_from_metadata(data) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index f25fa46197e..534b2db3e61 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -61,6 +61,7 @@ from litellm.proxy.common_utils.sse_keepalive import ( wrap_sse_stream_with_keepalive_pings, ) from litellm.proxy.dd_span_tagger import DDSpanTagger +from litellm.proxy.guardrails.auto_router_compression import arm_pre_call as _arm_auto_router_compression from litellm.proxy.route_llm_request import route_request from litellm.proxy.utils import ProxyLogging, _check_and_merge_model_level_guardrails from litellm.router import Router @@ -2004,6 +2005,12 @@ class ProxyBaseLLMRequestProcessing: trust_client_model_info=False, ) + # An auto router with its own compression policy is authoritative for this + # request: suppress every other compression guardrail and arm whichever one + # the policy names for the model call, before those guardrails get a chance + # to run below. + self.data = await _arm_auto_router_compression(data=self.data, llm_router=llm_router) + self.data = await proxy_logging_obj.pre_call_hook( user_api_key_dict=user_api_key_dict, data=self.data, diff --git a/litellm/proxy/guardrails/auto_router_compression.py b/litellm/proxy/guardrails/auto_router_compression.py new file mode 100644 index 00000000000..c3fba937d22 --- /dev/null +++ b/litellm/proxy/guardrails/auto_router_compression.py @@ -0,0 +1,211 @@ +""" +Decouples prompt compression between an auto router's routing decision and the +model it routes to. An auto router marker deployment may set +``auto_router_routing_compression`` and/or ``auto_router_model_compression`` in its +``litellm_params`` to name the compression guardrail that hop should use, or +``"none"`` to run no compression on that hop. Neither key set means the request's +own compression guardrails (key/team/model-level, or an "Always on" guardrail) +apply to both hops unchanged, exactly as before this feature existed. + +Once either key is set, this auto router is authoritative: every other compression +guardrail is suppressed for that request, and only these two settings decide what +each hop sees. +""" + +import copy +from collections.abc import Mapping +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Final + +from litellm._logging import verbose_proxy_logger +from litellm.constants import AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY +from litellm.litellm_core_utils.core_helpers import ( + get_metadata_variable_name_from_kwargs, + get_or_create_metadata_bucket, +) +from litellm.router_utils.auto_router_model_naming import AUTO_ROUTER_MODEL_PREFIX +from litellm.types.utils import GenericGuardrailAPIInputs + +if TYPE_CHECKING: + from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.router import Router +else: + CustomGuardrail = Any + Router = Any + +COMPRESSION_GUARDRAIL_PROVIDERS: Final = frozenset({"headroom", "compresr"}) +_NO_COMPRESSION: Final = "none" + +# Metadata key stashing the pre-compression messages so a routing decision that +# names a different compression than the model call still compresses the +# original text, not whatever the model-side guardrail already rewrote it to. +AUTO_ROUTER_ROUTING_MESSAGES_SNAPSHOT_KEY: Final = "_auto_router_routing_messages_snapshot" + + +@dataclass(frozen=True, slots=True) +class AutoRouterCompressionPolicy: + """An auto router's compression choice for each hop. ``None`` means no compression.""" + + routing: str | None + model: str | None + + @property + def is_same(self) -> bool: + return self.routing == self.model + + +def _normalized_compression_choice(raw: object) -> str | None: + if not isinstance(raw, str) or not raw: + return None + return None if raw.strip().lower() == _NO_COMPRESSION else raw + + +def policy_from_litellm_params(litellm_params: Mapping[str, object]) -> AutoRouterCompressionPolicy | None: + raw_routing: Final = litellm_params.get("auto_router_routing_compression") + raw_model: Final = litellm_params.get("auto_router_model_compression") + if raw_routing is None and raw_model is None: + return None + return AutoRouterCompressionPolicy( + routing=_normalized_compression_choice(raw_routing), + model=_normalized_compression_choice(raw_model), + ) + + +def policy_for_model( + llm_router: "Router | None", model_alias: str, team_id: str | None +) -> AutoRouterCompressionPolicy | None: + """The compression policy declared by the auto router marker deployment `model_alias` resolves to. + + Mirrors the alias lookup in ``_check_and_merge_model_level_guardrails``: this runs + before routing has picked a strategy, so it takes the first marker deployment for + the alias rather than disambiguating by request tags. + """ + if llm_router is None: + return None + deployments: Final = llm_router.get_model_list(model_name=model_alias, team_id=team_id) or [] + for deployment in deployments: + litellm_params: Final = deployment.get("litellm_params") or {} + model_field = litellm_params.get("model") + if not isinstance(model_field, str) or not model_field.startswith(AUTO_ROUTER_MODEL_PREFIX): + continue + policy = policy_from_litellm_params(litellm_params) + if policy is not None: + return policy + return None + + +def _active_compression_guardrail_names() -> frozenset[str]: + """Names of every currently-active guardrail whose type is a compression guardrail.""" + import litellm + from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.proxy.guardrails.guardrail_registry import guardrail_class_registry + + compression_classes: Final = tuple( + cls for name, cls in guardrail_class_registry.items() if name in COMPRESSION_GUARDRAIL_PROVIDERS + ) + if not compression_classes: + return frozenset() + active: Final = litellm.logging_callback_manager.get_custom_loggers_for_type(callback_type=CustomGuardrail) + return frozenset( + cb.guardrail_name for cb in active if isinstance(cb, compression_classes) and cb.guardrail_name + ) + + +async def arm_pre_call(data: dict, llm_router: "Router | None") -> dict: + """Apply an auto router's compression policy, if any, before guardrails run. + + Suppresses every other compression guardrail, re-enables the model-side + guardrail the policy names (if any) even when it isn't ``default_on``, and + snapshots the pre-compression messages so the routing decision can compress + them independently of whatever the model-side guardrail does to `data`. + """ + if llm_router is None: + return data + + model_alias: Final = data.get("model") + if not isinstance(model_alias, str) or not model_alias: + return data + + # Read-only until a policy is confirmed: creating the metadata bucket for every + # request, including the vast majority with no auto-router compression policy, + # would be an unwanted side effect of merely checking for one. + metadata_key: Final = get_metadata_variable_name_from_kwargs(data) + existing_bucket: Final = data.get(metadata_key) + other_bucket: Final = data.get("metadata" if metadata_key == "litellm_metadata" else "litellm_metadata") + team_id: Final = (existing_bucket.get("user_api_key_team_id") if isinstance(existing_bucket, dict) else None) or ( + other_bucket.get("user_api_key_team_id") if isinstance(other_bucket, dict) else None + ) + + policy: Final = policy_for_model(llm_router=llm_router, model_alias=model_alias, team_id=team_id) + if policy is None: + return data + + _, metadata = get_or_create_metadata_bucket(data) + suppressed: Final = _active_compression_guardrail_names() - ({policy.model} if policy.model else set()) + if suppressed: + metadata[AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY] = sorted(suppressed) + + if policy.model is not None: + requested = metadata.get("guardrails") + if isinstance(requested, list): + if policy.model not in requested: + requested.append(policy.model) + else: + metadata["guardrails"] = [policy.model] + + from litellm.litellm_core_utils.prompt_templates.factory import resolve_structured_messages + + snapshot: Final = resolve_structured_messages(messages=data.get("messages"), request_kwargs=data) + if snapshot is not None: + metadata[AUTO_ROUTER_ROUTING_MESSAGES_SNAPSHOT_KEY] = copy.deepcopy(snapshot) + + return data + + +async def messages_for_routing( + policy: AutoRouterCompressionPolicy | None, + messages: list[dict[str, Any]] | None, + request_kwargs: Mapping[str, object], +) -> list[dict[str, Any]] | None: + """Messages to use for a routing decision, compressed per `policy.routing`. + + Returns None when there is no policy or the policy's routing side names no + compression, meaning the caller should route on whatever messages it already + has. The model call is untouched by this function either way: model-side + compression, if any, already ran as an ordinary pre-call guardrail before the + router was ever reached. + """ + if policy is None or policy.routing is None: + return None + + from litellm.proxy.common_utils.registry_read_through import ( + get_initialized_guardrail_with_read_through, + ) + + metadata_key: Final = get_metadata_variable_name_from_kwargs(request_kwargs) + metadata: Final = request_kwargs.get(metadata_key) + snapshot: Final = metadata.get(AUTO_ROUTER_ROUTING_MESSAGES_SNAPSHOT_KEY) if isinstance(metadata, dict) else None + original: Final = snapshot if isinstance(snapshot, list) else messages + if not original: + return None + + guardrail: Final = await get_initialized_guardrail_with_read_through(policy.routing) + if guardrail is None: + verbose_proxy_logger.warning( + "AutoRouter compression: guardrail '%s' not found; routing on uncompressed messages", policy.routing + ) + return None + + inputs: GenericGuardrailAPIInputs = {"structured_messages": [dict(m) for m in original]} + # A throwaway request_data: apply_guardrail writes its stats onto this dict, not + # the real request's metadata, so routing-side compression never double-counts + # against extract_compression_saved_tokens's model-savings accounting. + throwaway_request_data: Final[dict[str, object]] = { + "messages": original, + "model": request_kwargs.get("model"), + } + result: Final = await guardrail.apply_guardrail( + inputs=inputs, request_data=throwaway_request_data, input_type="request" + ) + compressed = result.get("structured_messages") + return compressed if isinstance(compressed, list) else original diff --git a/litellm/router.py b/litellm/router.py index 6c7611c6236..8149ec60ddc 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -13037,13 +13037,46 @@ class Router: ) return None + from litellm.proxy.guardrails.auto_router_compression import ( + messages_for_routing, + policy_from_litellm_params, + ) + + marker_params: Final = self._alias_marker_litellm_params(registered_model_name, selected_strategy.tags) + compression_policy: Final = policy_from_litellm_params(marker_params) if marker_params else None + # When both hops share the same compression, the model-side guardrail already + # ran in the proxy's ordinary pre-call hook and compressed `messages` in place + # (arm_pre_call armed it whether or not it is `default_on`); reuse that result + # for routing too instead of paying for a second compression call against the + # same content. + needs_independent_routing_compression: Final = compression_policy is not None and not ( + compression_policy.is_same and compression_policy.model is not None + ) + routing_messages: Final = ( + await messages_for_routing(policy=compression_policy, messages=messages, request_kwargs=request_kwargs) + if needs_independent_routing_compression + else None + ) + pre_routing_hook_response: Final = await selected_strategy.strategy.async_pre_routing_hook( model=registered_model_name, request_kwargs=request_kwargs, - messages=messages, + messages=routing_messages if routing_messages is not None else messages, input=input, specific_deployment=specific_deployment, ) + # The strategy only echoes back whatever `messages` it was handed, so a + # routing-only compression must not leak into the response: the model call + # and downstream deployment-context filtering both key off this field. + # Compared by value, not identity: PreRoutingHookResponse is a pydantic model, + # and pydantic reconstructs a validated list field rather than keeping the + # exact object passed in, even when nothing about it changed. + if ( + pre_routing_hook_response is not None + and routing_messages is not None + and pre_routing_hook_response.messages == routing_messages + ): + pre_routing_hook_response = pre_routing_hook_response.model_copy(update={"messages": messages}) self._record_routing_decision( request_kwargs=request_kwargs, routing_decision=(pre_routing_hook_response.routing_decision if pre_routing_hook_response else None), @@ -13100,9 +13133,16 @@ class Router: return pre_routing_hook_response - def _forwardable_alias_marker_params( + def _alias_marker_litellm_params( self, model: str, strategy_tags: tuple[str, ...] - ) -> tuple[tuple[str, object], ...]: + ) -> Mapping[str, object] | None: + """The auto-router marker deployment's own `litellm_params` for `model`, tag-scoped. + + Shared by `_forwardable_alias_marker_params` (forwarding api_base/api_key/... + gaps onto the routed deployment) and the auto-router compression policy lookup + (reading `auto_router_routing_compression`/`auto_router_model_compression`), so + both read the same marker row when an alias has more than one, tag-scoped marker. + """ marker_params: Final = tuple( litellm_params for idx in self.model_name_to_deployment_indices.get(model, ()) @@ -13112,7 +13152,12 @@ class Router: tag_matched: Final = tuple( params for params in marker_params if tuple(params.get("tags") or ()) == strategy_tags ) - selected: Final = tag_matched[0] if tag_matched else (marker_params[0] if marker_params else None) + return tag_matched[0] if tag_matched else (marker_params[0] if marker_params else None) + + def _forwardable_alias_marker_params( + self, model: str, strategy_tags: tuple[str, ...] + ) -> tuple[tuple[str, object], ...]: + selected: Final = self._alias_marker_litellm_params(model, strategy_tags) if selected is None: return () return tuple( diff --git a/litellm/types/router.py b/litellm/types/router.py index 7ebd50f1328..f5295d6569c 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -359,6 +359,10 @@ class GenericLiteLLMParams(CredentialLiteLLMParams, CustomPricingLiteLLMParams): auto_router_default_model: str | None = None auto_router_embedding_model: str | None = None auto_router_max_input_chars: int | None = None + # Compression policy for the two hops of a routed request. Both unset means the + # request's own compression guardrails apply to both, as they always have. + auto_router_routing_compression: str | None = None + auto_router_model_compression: str | None = None # complexity-router params complexity_router_config: dict | None = None diff --git a/litellm/types/utils.py b/litellm/types/utils.py index c1d90694f14..a1fda3d0524 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3713,6 +3713,8 @@ all_litellm_params = ( "auto_router_default_model", "auto_router_embedding_model", "auto_router_max_input_chars", + "auto_router_routing_compression", + "auto_router_model_compression", "complexity_router_config", "complexity_router_default_model", "adaptive_router_config", diff --git a/tests/test_litellm/integrations/test_custom_guardrail.py b/tests/test_litellm/integrations/test_custom_guardrail.py index 6edc2c9bf77..1fb4299cb56 100644 --- a/tests/test_litellm/integrations/test_custom_guardrail.py +++ b/tests/test_litellm/integrations/test_custom_guardrail.py @@ -518,6 +518,50 @@ class TestCustomGuardrailShouldRunGuardrail: is True ) + def test_should_run_guardrail_suppressed_by_auto_router_compression(self): + """An auto router's own compression policy can suppress an otherwise-eligible + guardrail, even one that is default_on and explicitly requested.""" + from litellm.constants import AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY + from litellm.types.guardrails import GuardrailEventHooks + + always_on = CustomGuardrail( + guardrail_name="headroom-default", + default_on=True, + event_hook=GuardrailEventHooks.pre_call, + ) + data = { + "model": "smart-router", + "metadata": { + AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY: ["headroom-default"], + }, + } + + assert ( + always_on.should_run_guardrail(data=data, event_type=GuardrailEventHooks.pre_call) + is False + ) + + def test_should_run_guardrail_suppression_list_does_not_affect_other_names(self): + from litellm.constants import AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY + from litellm.types.guardrails import GuardrailEventHooks + + always_on = CustomGuardrail( + guardrail_name="headroom-default", + default_on=True, + event_hook=GuardrailEventHooks.pre_call, + ) + data = { + "model": "smart-router", + "metadata": { + AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY: ["some-other-guardrail"], + }, + } + + assert ( + always_on.should_run_guardrail(data=data, event_type=GuardrailEventHooks.pre_call) + is True + ) + class TestApplyGuardrailCheck: def test_apply_guardrail_check_only_on_direct_implementation(self): diff --git a/tests/test_litellm/proxy/guardrails/test_auto_router_compression.py b/tests/test_litellm/proxy/guardrails/test_auto_router_compression.py new file mode 100644 index 00000000000..b2e83e75768 --- /dev/null +++ b/tests/test_litellm/proxy/guardrails/test_auto_router_compression.py @@ -0,0 +1,285 @@ +""" +Unit tests for litellm.proxy.guardrails.auto_router_compression. + +Covers: +- policy_from_litellm_params: absent keys mean no policy; the "none" sentinel + normalizes to explicit no-compression within an active policy; is_same +- policy_for_model: finds the auto-router marker deployment for an alias +- arm_pre_call: no-op without a policy; suppresses active compression guardrails; + arms the model-side guardrail even when it isn't default_on; snapshots messages +- messages_for_routing: no-op without a policy or an unset routing side; compresses + via the named guardrail's apply_guardrail; never writes stats onto the caller's + own request_kwargs (regression for double-counted compression savings) +""" + +from typing import Any + +import pytest + +from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.proxy.guardrails.auto_router_compression import ( + AUTO_ROUTER_ROUTING_MESSAGES_SNAPSHOT_KEY, + AutoRouterCompressionPolicy, + arm_pre_call, + messages_for_routing, + policy_for_model, + policy_from_litellm_params, +) +from litellm.constants import AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY +from litellm.types.utils import GenericGuardrailAPIInputs + + +class TestPolicyFromLitellmParams: + def test_neither_key_set_is_no_policy(self): + assert policy_from_litellm_params({}) is None + + def test_routing_only(self): + policy = policy_from_litellm_params({"auto_router_routing_compression": "headroom-a"}) + assert policy == AutoRouterCompressionPolicy(routing="headroom-a", model=None) + + def test_none_sentinel_normalizes_to_no_compression(self): + policy = policy_from_litellm_params( + {"auto_router_routing_compression": "headroom-a", "auto_router_model_compression": "none"} + ) + assert policy == AutoRouterCompressionPolicy(routing="headroom-a", model=None) + + def test_none_sentinel_is_case_insensitive(self): + policy = policy_from_litellm_params({"auto_router_routing_compression": "NONE"}) + assert policy == AutoRouterCompressionPolicy(routing=None, model=None) + + def test_is_same_true_for_matching_names(self): + policy = policy_from_litellm_params( + {"auto_router_routing_compression": "x", "auto_router_model_compression": "x"} + ) + assert policy.is_same is True + + def test_is_same_false_for_different_names(self): + policy = policy_from_litellm_params( + {"auto_router_routing_compression": "x", "auto_router_model_compression": "y"} + ) + assert policy.is_same is False + + def test_is_same_true_when_both_no_compression(self): + policy = policy_from_litellm_params( + {"auto_router_routing_compression": "none", "auto_router_model_compression": "none"} + ) + assert policy.is_same is True + + +class _FakeRouter: + """Minimal stand-in for litellm.Router.get_model_list, for policy_for_model.""" + + def __init__(self, deployments: list[dict[str, Any]]): + self._deployments = deployments + + def get_model_list(self, model_name, team_id=None): + return [d for d in self._deployments if d.get("model_name") == model_name] + + +class TestPolicyForModel: + def test_no_router_returns_none(self): + assert policy_for_model(llm_router=None, model_alias="smart-router", team_id=None) is None + + def test_no_marker_deployment_returns_none(self): + router = _FakeRouter( + [{"model_name": "smart-router", "litellm_params": {"model": "openai/gpt-4o-mini"}}] + ) + assert policy_for_model(llm_router=router, model_alias="smart-router", team_id=None) is None + + def test_marker_deployment_without_policy_returns_none(self): + router = _FakeRouter( + [{"model_name": "smart-router", "litellm_params": {"model": "auto_router/complexity_router"}}] + ) + assert policy_for_model(llm_router=router, model_alias="smart-router", team_id=None) is None + + def test_marker_deployment_with_policy_is_found(self): + router = _FakeRouter( + [ + { + "model_name": "smart-router", + "litellm_params": { + "model": "auto_router/complexity_router", + "auto_router_routing_compression": "headroom-a", + "auto_router_model_compression": "none", + }, + } + ] + ) + policy = policy_for_model(llm_router=router, model_alias="smart-router", team_id=None) + assert policy == AutoRouterCompressionPolicy(routing="headroom-a", model=None) + + +class _RecordingCompressionGuardrail(CustomGuardrail): + """A guardrail whose apply_guardrail marks every text message as compressed.""" + + def __init__(self, guardrail_name: str): + super().__init__(guardrail_name=guardrail_name) + self.request_data_seen: list[dict] = [] + + async def apply_guardrail( + self, inputs: GenericGuardrailAPIInputs, request_data: dict, input_type: str, logging_obj=None + ) -> GenericGuardrailAPIInputs: + self.request_data_seen.append(request_data) + structured_messages = inputs.get("structured_messages") or [] + compressed = [ + {**m, "content": f"[COMPRESSED] {m.get('content')}"} for m in structured_messages + ] + return {**inputs, "structured_messages": compressed} + + +@pytest.fixture +def registered_guardrail(): + import litellm + + guardrail = _RecordingCompressionGuardrail(guardrail_name="fake-compress") + litellm.logging_callback_manager.add_litellm_callback(guardrail) + yield guardrail + litellm.logging_callback_manager.remove_callback_from_all_lists(guardrail) + + +class TestArmPreCall: + @pytest.mark.asyncio + async def test_no_router_is_noop(self): + data = {"model": "smart-router", "messages": [{"role": "user", "content": "hi"}]} + result = await arm_pre_call(data=data, llm_router=None) + assert result == data + assert "metadata" not in result + + @pytest.mark.asyncio + async def test_no_policy_does_not_create_metadata_bucket(self): + router = _FakeRouter( + [{"model_name": "smart-router", "litellm_params": {"model": "openai/gpt-4o-mini"}}] + ) + data = {"model": "smart-router", "messages": [{"role": "user", "content": "hi"}]} + result = await arm_pre_call(data=data, llm_router=router) + assert "metadata" not in result + assert "litellm_metadata" not in result + + @pytest.mark.asyncio + async def test_policy_suppresses_active_compression_guardrails(self, monkeypatch): + from litellm.proxy.guardrails import guardrail_registry + + monkeypatch.setitem( + guardrail_registry.guardrail_class_registry, "fake-provider", _RecordingCompressionGuardrail + ) + monkeypatch.setattr( + "litellm.proxy.guardrails.auto_router_compression.COMPRESSION_GUARDRAIL_PROVIDERS", + frozenset({"fake-provider"}), + ) + import litellm + + always_on = _RecordingCompressionGuardrail(guardrail_name="always-on-compression") + litellm.logging_callback_manager.add_litellm_callback(always_on) + try: + router = _FakeRouter( + [ + { + "model_name": "smart-router", + "litellm_params": { + "model": "auto_router/complexity_router", + "auto_router_routing_compression": "headroom-a", + "auto_router_model_compression": "none", + }, + } + ] + ) + data = {"model": "smart-router", "messages": [{"role": "user", "content": "hi"}]} + result = await arm_pre_call(data=data, llm_router=router) + suppressed = result["metadata"][AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY] + assert "always-on-compression" in suppressed + finally: + litellm.logging_callback_manager.remove_callback_from_all_lists(always_on) + + @pytest.mark.asyncio + async def test_model_side_guardrail_is_requested_even_when_not_default_on(self): + router = _FakeRouter( + [ + { + "model_name": "smart-router", + "litellm_params": { + "model": "auto_router/complexity_router", + "auto_router_routing_compression": "none", + "auto_router_model_compression": "headroom-b", + }, + } + ] + ) + data = {"model": "smart-router", "messages": [{"role": "user", "content": "hi"}]} + result = await arm_pre_call(data=data, llm_router=router) + assert result["metadata"]["guardrails"] == ["headroom-b"] + + @pytest.mark.asyncio + async def test_snapshots_original_messages(self): + router = _FakeRouter( + [ + { + "model_name": "smart-router", + "litellm_params": { + "model": "auto_router/complexity_router", + "auto_router_routing_compression": "headroom-a", + "auto_router_model_compression": "none", + }, + } + ] + ) + original_messages = [{"role": "user", "content": "hi"}] + data = {"model": "smart-router", "messages": original_messages} + result = await arm_pre_call(data=data, llm_router=router) + snapshot = result["metadata"][AUTO_ROUTER_ROUTING_MESSAGES_SNAPSHOT_KEY] + assert snapshot == original_messages + assert snapshot is not original_messages # a copy, not the live reference + + +class TestMessagesForRouting: + @pytest.mark.asyncio + async def test_no_policy_returns_none(self): + assert await messages_for_routing(policy=None, messages=[], request_kwargs={}) is None + + @pytest.mark.asyncio + async def test_routing_side_unset_returns_none(self): + policy = AutoRouterCompressionPolicy(routing=None, model="headroom-a") + assert await messages_for_routing(policy=policy, messages=[], request_kwargs={}) is None + + @pytest.mark.asyncio + async def test_unknown_guardrail_name_returns_none(self): + policy = AutoRouterCompressionPolicy(routing="does-not-exist", model=None) + messages = [{"role": "user", "content": "hi"}] + result = await messages_for_routing(policy=policy, messages=messages, request_kwargs={}) + assert result is None + + @pytest.mark.asyncio + async def test_compresses_via_the_named_guardrail(self, registered_guardrail): + policy = AutoRouterCompressionPolicy(routing="fake-compress", model=None) + messages = [{"role": "user", "content": "hello world"}] + result = await messages_for_routing(policy=policy, messages=messages, request_kwargs={}) + assert result == [{"role": "user", "content": "[COMPRESSED] hello world"}] + + @pytest.mark.asyncio + async def test_uses_the_snapshot_when_present(self, registered_guardrail): + policy = AutoRouterCompressionPolicy(routing="fake-compress", model=None) + snapshot = [{"role": "user", "content": "original"}] + request_kwargs = {"metadata": {AUTO_ROUTER_ROUTING_MESSAGES_SNAPSHOT_KEY: snapshot}} + # `messages` here stands in for whatever a model-side guardrail already + # rewrote `data["messages"]` to -- routing must ignore it and compress the + # pristine snapshot instead. + already_rewritten = [{"role": "user", "content": "rewritten by another guardrail"}] + result = await messages_for_routing( + policy=policy, messages=already_rewritten, request_kwargs=request_kwargs + ) + assert result == [{"role": "user", "content": "[COMPRESSED] original"}] + + @pytest.mark.asyncio + async def test_guardrail_receives_a_throwaway_request_data_not_the_real_request_kwargs( + self, registered_guardrail + ): + """Regression: a real compression guardrail writes its stats onto whatever + `request_data` dict it's given (`add_standard_logging_guardrail_information_to_ + request_data`). If that were the caller's own `request_kwargs`, routing-side + compression would double-count into extract_compression_saved_tokens, which + sums every guardrail_information entry on the real request's metadata.""" + policy = AutoRouterCompressionPolicy(routing="fake-compress", model=None) + messages = [{"role": "user", "content": "hi"}] + request_kwargs = {"metadata": {}} + await messages_for_routing(policy=policy, messages=messages, request_kwargs=request_kwargs) + assert registered_guardrail.request_data_seen[0] is not request_kwargs + assert request_kwargs == {"metadata": {}} diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index f7fe6ad9d39..c0809e53d2e 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -376,6 +376,60 @@ class TestProxyBaseLLMRequestProcessing: assert "litellm_logging_obj" not in persisted_body json.dumps(persisted_body) + @pytest.mark.asyncio + async def test_common_processing_pre_call_logic_arms_auto_router_compression_before_guardrails( + self, monkeypatch + ): + """arm_pre_call must run before pre_call_hook: an auto router's own compression + policy has to be in `data["metadata"]` (naming the model-side guardrail so it + runs even if it isn't default_on) by the time guardrails see the request.""" + processing_obj = ProxyBaseLLMRequestProcessing(data={}) + mock_request = MagicMock(spec=Request) + mock_request.headers = {} + + async def mock_add_litellm_data_to_request(*args, **kwargs): + return {"model": "smart-router", "messages": [{"role": "user", "content": "hi"}]} + + seen_metadata: dict = {} + + async def mock_pre_call_hook(user_api_key_dict, data, call_type): + seen_metadata.update(data.get("metadata") or {}) + return data + + mock_proxy_logging_obj = MagicMock(spec=ProxyLogging) + mock_proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=mock_pre_call_hook) + monkeypatch.setattr( + litellm.proxy.common_request_processing, + "add_litellm_data_to_request", + mock_add_litellm_data_to_request, + ) + + fake_llm_router = MagicMock() + fake_llm_router.get_model_list.return_value = [ + { + "model_name": "smart-router", + "litellm_params": { + "model": "auto_router/complexity_router", + "auto_router_routing_compression": "none", + "auto_router_model_compression": "headroom-model", + }, + } + ] + mock_proxy_config = MagicMock(spec=ProxyConfig) + mock_proxy_config._get_hierarchical_router_settings = AsyncMock(return_value=None) + + await processing_obj.common_processing_pre_call_logic( + request=mock_request, + general_settings={}, + user_api_key_dict=ProxyUserAPIKeyAuth(api_key="sk-test"), + proxy_logging_obj=mock_proxy_logging_obj, + proxy_config=mock_proxy_config, + route_type="acompletion", + llm_router=fake_llm_router, + ) + + assert seen_metadata.get("guardrails") == ["headroom-model"] + def test_add_dd_apm_tags_for_litellm_call_id_uses_dd_tracing_helper(self, monkeypatch): mock_set_active_span_tag = MagicMock(return_value=True) import litellm.proxy.dd_span_tagger diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 5fc96bcfbb1..cdae3b131ae 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -18,6 +18,7 @@ import pytest import litellm from litellm import Router from litellm.exceptions import MidStreamFallbackError +from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.integrations.custom_logger import CustomLogger from litellm.llms.bedrock.common_utils import BedrockError from litellm.llms.anthropic.experimental_pass_through.messages.agentic_streaming_iterator import ( @@ -9995,6 +9996,156 @@ class TestModelGroupAliasReachesPreRoutingStrategies: ) +class TestAutoRouterCompressionDecoupling: + """An auto router's `auto_router_routing_compression` / `auto_router_model_compression` + decouple what the routing decision sees from what the model call sees. The one + assertion that must hold under any mutation: the strategy can be routed on + compressed text while the caller's own `messages` list - the one that would reach + the model - is never touched.""" + + class _RecordingStrategy: + """Echoes back whatever `messages` it was handed, like every real strategy does.""" + + def __init__(self): + self.received_messages: list[dict] | None = None + + async def async_pre_routing_hook( + self, model, request_kwargs, messages=None, input=None, specific_deployment=False + ): + from litellm.types.router import PreRoutingHookResponse + + self.received_messages = messages + return PreRoutingHookResponse(model="gemini-flash", messages=messages) + + class _CompressingGuardrail(CustomGuardrail): + def __init__(self, guardrail_name: str): + super().__init__(guardrail_name=guardrail_name) + self.call_count = 0 + + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + self.call_count += 1 + structured_messages = inputs.get("structured_messages") or [] + compressed = [ + {**m, "content": f"[COMPRESSED] {m.get('content')}"} for m in structured_messages + ] + return {**inputs, "structured_messages": compressed} + + @staticmethod + def _messages() -> list[dict[str, str]]: + return [{"role": "user", "content": "What is the capital of France?"}] + + def _router(self, marker_litellm_params: dict) -> tuple[litellm.Router, "_RecordingStrategy"]: + from litellm.types.router import TaggedPreRoutingStrategy + + tiers = dict.fromkeys(("SIMPLE", "MEDIUM", "COMPLEX", "REASONING"), "gemini-flash") + router = litellm.Router( + model_list=[ + { + "model_name": "smart-router", + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": {"tiers": tiers}, + "complexity_router_default_model": "gemini-flash", + **marker_litellm_params, + }, + }, + { + "model_name": "gemini-flash", + "litellm_params": {"model": "gemini/gemini-3.6-flash", "mock_response": "routed by the tier"}, + }, + ], + ) + for name in ("auto_routers", "complexity_routers", "adaptive_routers", "quality_routers"): + setattr(router, name, {}) + strategy = self._RecordingStrategy() + router.complexity_routers = {"smart-router": [TaggedPreRoutingStrategy(tags=(), strategy=strategy)]} + return router, strategy + + @pytest.fixture + def registered_guardrail(self): + guardrail = self._CompressingGuardrail(guardrail_name="fake-compress") + litellm.logging_callback_manager.add_litellm_callback(guardrail) + yield guardrail + litellm.logging_callback_manager.remove_callback_from_all_lists(guardrail) + + @pytest.mark.asyncio + async def test_routing_side_compression_never_reaches_the_caller_messages(self, registered_guardrail): + router, strategy = self._router( + { + "auto_router_routing_compression": "fake-compress", + "auto_router_model_compression": "none", + } + ) + original_messages = self._messages() + + response = await router.async_pre_routing_hook( + model="smart-router", request_kwargs={"metadata": {}}, messages=original_messages + ) + + assert strategy.received_messages == [ + {"role": "user", "content": "[COMPRESSED] What is the capital of France?"} + ] + assert response.messages == original_messages + + @pytest.mark.asyncio + async def test_model_side_compression_alone_leaves_routing_uncompressed(self, registered_guardrail): + router, strategy = self._router( + { + "auto_router_routing_compression": "none", + "auto_router_model_compression": "fake-compress", + } + ) + original_messages = self._messages() + + response = await router.async_pre_routing_hook( + model="smart-router", request_kwargs={"metadata": {}}, messages=original_messages + ) + + assert strategy.received_messages == original_messages + assert response.messages == original_messages + assert registered_guardrail.call_count == 0 + + @pytest.mark.asyncio + async def test_same_compression_on_both_hops_compresses_once(self, registered_guardrail): + """The same/different distinction exists so a shared choice does not pay for + compression twice: by the time the router runs, `messages` already reflects + whatever the ordinary pre-call guardrail pipeline did for the model call, so + the routing decision must reuse it rather than calling the guardrail again.""" + router, strategy = self._router( + { + "auto_router_routing_compression": "fake-compress", + "auto_router_model_compression": "fake-compress", + } + ) + # Stands in for what the proxy's ordinary pre-call guardrail pipeline would + # have already produced for the model call, since `auto_router_model_compression` + # names a guardrail: the router never triggers that pipeline itself. + already_compressed_messages = [ + {"role": "user", "content": "[COMPRESSED] What is the capital of France?"} + ] + + response = await router.async_pre_routing_hook( + model="smart-router", request_kwargs={"metadata": {}}, messages=already_compressed_messages + ) + + assert strategy.received_messages == already_compressed_messages + assert response.messages == already_compressed_messages + assert registered_guardrail.call_count == 0 + + @pytest.mark.asyncio + async def test_no_policy_is_fully_unaffected(self, registered_guardrail): + router, strategy = self._router({}) + original_messages = self._messages() + + response = await router.async_pre_routing_hook( + model="smart-router", request_kwargs={"metadata": {}}, messages=original_messages + ) + + assert strategy.received_messages is original_messages + assert response.messages == original_messages + assert registered_guardrail.call_count == 0 + + @pytest.mark.usefixtures("local_model_cost_map") class TestAzureBaseModelFallbackLogging: """When an azure deployment has no base_model but its model name is a known diff --git a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx index 2a024ab7fdf..42115265034 100644 --- a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx @@ -48,6 +48,8 @@ import EscalationKeywords from "./EscalationKeywords"; import KeywordTierRules, { KeywordTierRule } from "./KeywordTierRules"; import SemanticKeywordMatching from "./SemanticKeywordMatching"; import { type DimensionWeights, type TierBoundaries, type TokenThresholds } from "./heuristic_scoring_knobs"; +import CompressionControls from "./CompressionControls"; +import { type AutoRouterCompressionState, DEFAULT_AUTO_ROUTER_COMPRESSION } from "./buildAutoRouterCompression"; export type { DimensionWeights, TierBoundaries, TokenThresholds }; export type { CustomTierSet, TierRow } from "./tier_rows"; @@ -490,6 +492,10 @@ interface ComplexityRouterConfigProps { onMatchThresholdChange?: (threshold: number) => void; escalationKeywords?: string[]; onEscalationKeywordsChange?: (keywords: string[]) => void; + // Optional: not part of complexity_router_config, since it applies to every + // pre-routing strategy, not just the complexity router. + autoRouterCompression?: AutoRouterCompressionState; + onAutoRouterCompressionChange?: (state: AutoRouterCompressionState) => void; showValidationErrors?: boolean; } @@ -611,6 +617,8 @@ const ComplexityRouterConfig: React.FC = ({ onMatchThresholdChange = () => {}, escalationKeywords = [], onEscalationKeywordsChange, + autoRouterCompression = DEFAULT_AUTO_ROUTER_COMPRESSION, + onAutoRouterCompressionChange, showValidationErrors = false, }) => { const customTierSet = value.custom_tier_set; @@ -875,6 +883,32 @@ const ComplexityRouterConfig: React.FC = ({ }, ] : []), + ...(onAutoRouterCompressionChange + ? [ + { + key: "compression", + label: Advanced: Compression, + children: ( + + onAutoRouterCompressionChange({ + ...autoRouterCompression, + routing, + sameAsRouting: routing === undefined ? true : autoRouterCompression.sameAsRouting, + }) + } + sameAsRouting={autoRouterCompression.sameAsRouting} + onSameAsRoutingChange={(sameAsRouting) => + onAutoRouterCompressionChange({ ...autoRouterCompression, sameAsRouting }) + } + model={autoRouterCompression.model} + onModelChange={(model) => onAutoRouterCompressionChange({ ...autoRouterCompression, model })} + /> + ), + }, + ] + : []), ...(onKeywordTierRulesChange || onSemanticMatchingEnabledChange ? [ { diff --git a/ui/litellm-dashboard/src/components/add_model/CompressionControls.tsx b/ui/litellm-dashboard/src/components/add_model/CompressionControls.tsx new file mode 100644 index 00000000000..a0a240f76b0 --- /dev/null +++ b/ui/litellm-dashboard/src/components/add_model/CompressionControls.tsx @@ -0,0 +1,93 @@ +import { SimpleTooltip } from "@/components/ui/tooltip"; +import { SearchSelect, SearchSelectOption } from "@/components/shared/SearchSelect"; +import { Label } from "@/components/ui/label"; +import { RadioGroup, RadioGroupItem } from "@/components/ui/radio-group"; +import { Info } from "lucide-react"; +import React from "react"; +import { useGuardrails } from "@/app/(dashboard)/hooks/guardrails/useGuardrails"; +import { COMPRESSION_GUARDRAIL_PROVIDER } from "@/app/(dashboard)/cost-optimization/_components/helpers"; +import { NO_COMPRESSION } from "./buildAutoRouterCompression"; + +interface CompressionControlsProps { + routing: string | undefined; + onRoutingChange: (value: string | undefined) => void; + sameAsRouting: boolean; + onSameAsRoutingChange: (same: boolean) => void; + model: string | undefined; + onModelChange: (value: string | undefined) => void; +} + +const NONE_OPTION: SearchSelectOption = { label: "None (no compression)", value: NO_COMPRESSION }; + +const CompressionControls: React.FC = ({ + routing, + onRoutingChange, + sameAsRouting, + onSameAsRoutingChange, + model, + onModelChange, +}) => { + const { data } = useGuardrails(); + const compressionOptions: SearchSelectOption[] = (data?.guardrails ?? []) + .filter((g) => (g.litellm_params?.guardrail ?? "").toString().toLowerCase() === COMPRESSION_GUARDRAIL_PROVIDER) + .map((g) => ({ label: g.guardrail_name, value: g.guardrail_name })); + const options: SearchSelectOption[] = [NONE_OPTION, ...compressionOptions]; + + return ( +
+
+
+ Routing decision + + + +
+ onRoutingChange(value === "" ? undefined : value)} + placeholder="Inherit from the request's own compression guardrails" + emptyText="No compression guardrails found" + aria-label="Routing decision compression" + /> +
+ + {routing !== undefined && ( +
+ Model call + onSameAsRoutingChange(value === "same")} + className="w-full" + > +
+ + +
+ + + {!sameAsRouting && ( +
+ onModelChange(value === "" ? undefined : value)} + placeholder="None (no compression)" + emptyText="No compression guardrails found" + aria-label="Model call compression" + /> +
+ )} +
+ )} +
+ ); +}; + +export default CompressionControls; diff --git a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx index d2f6b10c3a6..5605a993ded 100644 --- a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx @@ -1,4 +1,12 @@ -import { renderWithProviders, screen, waitFor, within, fireEvent, testQueryClient } from "../../../tests/test-utils"; +import { + renderWithProviders, + screen, + waitFor, + within, + fireEvent, + testQueryClient, + chooseSelectOption, +} from "../../../tests/test-utils"; import userEvent from "@testing-library/user-event"; import { vi } from "vitest"; import AddAutoRouterTab from "./add_auto_router_tab"; @@ -522,6 +530,71 @@ describe("AddAutoRouterTab", () => { ); }); + describe("prompt compression", () => { + it("leaves both compression keys out of the create payload when the section is untouched", async () => { + const user = userEvent.setup(); + vi.mocked(getMissingTiersError).mockReturnValue(null); + + renderWithProviders(); + + await user.type(screen.getByPlaceholderText(/smart_router/i), "no-compression-router"); + await user.click(screen.getByRole("button", { name: /add auto router/i })); + + await waitFor(() => expect(handleAddAutoRouterSubmit).toHaveBeenCalled()); + const submitted = vi.mocked(handleAddAutoRouterSubmit).mock.calls.at(-1)?.[0]; + expect(submitted).not.toHaveProperty("auto_router_routing_compression"); + expect(submitted).not.toHaveProperty("auto_router_model_compression"); + }); + + it("mirrors an explicit no-compression routing choice onto the model call by default", async () => { + const user = userEvent.setup(); + vi.mocked(getMissingTiersError).mockReturnValue(null); + + renderWithProviders(); + + await user.type(screen.getByPlaceholderText(/smart_router/i), "no-compression-explicit-router"); + expandDetailedConfiguration(); + await user.click(screen.getByText("Advanced: Compression")); + await chooseSelectOption( + user, + screen.getByRole("combobox", { name: "Routing decision compression" }), + "None (no compression)", + ); + + await user.click(screen.getByRole("button", { name: /add auto router/i })); + + await waitFor(() => expect(handleAddAutoRouterSubmit).toHaveBeenCalled()); + const submitted = vi.mocked(handleAddAutoRouterSubmit).mock.calls.at(-1)?.[0]; + expect(submitted?.auto_router_routing_compression).toBe("none"); + expect(submitted?.auto_router_model_compression).toBe("none"); + }); + + it("defaults the model call to none when different is chosen but nothing is picked there", async () => { + const user = userEvent.setup(); + vi.mocked(getMissingTiersError).mockReturnValue(null); + + renderWithProviders(); + + await user.type(screen.getByPlaceholderText(/smart_router/i), "different-compression-router"); + expandDetailedConfiguration(); + await user.click(screen.getByText("Advanced: Compression")); + await chooseSelectOption( + user, + screen.getByRole("combobox", { name: "Routing decision compression" }), + "None (no compression)", + ); + await user.click(screen.getByText("Use a different compression")); + expect(screen.getByRole("combobox", { name: "Model call compression" })).toBeInTheDocument(); + + await user.click(screen.getByRole("button", { name: /add auto router/i })); + + await waitFor(() => expect(handleAddAutoRouterSubmit).toHaveBeenCalled()); + const submitted = vi.mocked(handleAddAutoRouterSubmit).mock.calls.at(-1)?.[0]; + expect(submitted?.auto_router_routing_compression).toBe("none"); + expect(submitted?.auto_router_model_compression).toBe("none"); + }); + }); + // The scalar floor is the one scorer knob with no group dict behind it, so its wiring into the create // payload is only proven end to end. 0 is the case a truthy check would silently drop. it("carries a reasoning override floor of 0 through to the create payload", async () => { diff --git a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx index a548c2c6533..decacac6501 100644 --- a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx +++ b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx @@ -32,6 +32,11 @@ import ComplexityRouterConfig, { } from "./ComplexityRouterConfig"; import { KeywordTierRule } from "./KeywordTierRules"; import { DEFAULT_ESCALATION_KEYWORDS } from "./EscalationKeywords"; +import { + type AutoRouterCompressionState, + buildAutoRouterCompressionParams, + DEFAULT_AUTO_ROUTER_COMPRESSION, +} from "./buildAutoRouterCompression"; import { DEFAULT_MATCH_THRESHOLD } from "./SemanticKeywordMatching"; import { BuildComplexityRouterConfigParams, @@ -194,6 +199,9 @@ const AddAutoRouterTab: React.FC = ({ const [embeddingModel, setEmbeddingModel] = useState(undefined); const [matchThreshold, setMatchThreshold] = useState(DEFAULT_MATCH_THRESHOLD); const [escalationKeywords, setEscalationKeywords] = useState(DEFAULT_ESCALATION_KEYWORDS); + const [autoRouterCompression, setAutoRouterCompression] = useState( + DEFAULT_AUTO_ROUTER_COMPRESSION, + ); const [showValidationErrors, setShowValidationErrors] = useState(false); const [editingTiers, setEditingTiers] = useState(false); const [isSubmitting, setIsSubmitting] = useState(false); @@ -461,6 +469,7 @@ const AddAutoRouterTab: React.FC = ({ model_type: "complexity_router", complexity_router_config: complexityRouterConfigPayload, model_access_group: form.getValues("model_access_group"), + ...buildAutoRouterCompressionParams(autoRouterCompression), }; await handleAddAutoRouterSubmit(submitValues, accessToken, () => form.reset(EMPTY_FORM_VALUES), handleOk); @@ -666,6 +675,8 @@ const AddAutoRouterTab: React.FC = ({ onMatchThresholdChange={setMatchThreshold} escalationKeywords={escalationKeywords} onEscalationKeywordsChange={setEscalationKeywords} + autoRouterCompression={autoRouterCompression} + onAutoRouterCompressionChange={setAutoRouterCompression} showValidationErrors={showValidationErrors} /> diff --git a/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.test.ts b/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.test.ts new file mode 100644 index 00000000000..b917fcedaa2 --- /dev/null +++ b/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.test.ts @@ -0,0 +1,93 @@ +import { + buildAutoRouterCompressionParams, + DEFAULT_AUTO_ROUTER_COMPRESSION, + hydrateAutoRouterCompression, + NO_COMPRESSION, +} from "./buildAutoRouterCompression"; + +describe("buildAutoRouterCompressionParams", () => { + it("omits both keys when routing was never configured", () => { + expect(buildAutoRouterCompressionParams(DEFAULT_AUTO_ROUTER_COMPRESSION)).toEqual({}); + }); + + it("mirrors routing onto model when same-as-routing is chosen", () => { + const params = buildAutoRouterCompressionParams({ + routing: "headroom-a", + sameAsRouting: true, + model: undefined, + }); + expect(params).toEqual({ + auto_router_routing_compression: "headroom-a", + auto_router_model_compression: "headroom-a", + }); + }); + + it("uses the explicit model choice when different is chosen", () => { + const params = buildAutoRouterCompressionParams({ + routing: "headroom-a", + sameAsRouting: false, + model: "headroom-b", + }); + expect(params).toEqual({ + auto_router_routing_compression: "headroom-a", + auto_router_model_compression: "headroom-b", + }); + }); + + it("defaults the model side to none when different is chosen but nothing is picked", () => { + const params = buildAutoRouterCompressionParams({ + routing: "headroom-a", + sameAsRouting: false, + model: undefined, + }); + expect(params).toEqual({ + auto_router_routing_compression: "headroom-a", + auto_router_model_compression: NO_COMPRESSION, + }); + }); + + it("sends the none sentinel when routing itself is explicitly turned off", () => { + const params = buildAutoRouterCompressionParams({ + routing: NO_COMPRESSION, + sameAsRouting: true, + model: undefined, + }); + expect(params).toEqual({ + auto_router_routing_compression: NO_COMPRESSION, + auto_router_model_compression: NO_COMPRESSION, + }); + }); +}); + +describe("hydrateAutoRouterCompression", () => { + it("returns the default state when neither key is set", () => { + expect(hydrateAutoRouterCompression({})).toEqual(DEFAULT_AUTO_ROUTER_COMPRESSION); + }); + + it("is same-as-routing when the model value matches routing", () => { + const state = hydrateAutoRouterCompression({ + auto_router_routing_compression: "headroom-a", + auto_router_model_compression: "headroom-a", + }); + expect(state).toEqual({ routing: "headroom-a", sameAsRouting: true, model: undefined }); + }); + + it("is different when the model value diverges from routing", () => { + const state = hydrateAutoRouterCompression({ + auto_router_routing_compression: "headroom-a", + auto_router_model_compression: "headroom-b", + }); + expect(state).toEqual({ routing: "headroom-a", sameAsRouting: false, model: "headroom-b" }); + }); + + it("treats a missing model key as same-as-routing", () => { + const state = hydrateAutoRouterCompression({ auto_router_routing_compression: "headroom-a" }); + expect(state).toEqual({ routing: "headroom-a", sameAsRouting: true, model: undefined }); + }); + + it("round-trips through buildAutoRouterCompressionParams", () => { + const original = { auto_router_routing_compression: "headroom-a", auto_router_model_compression: "none" }; + const rebuilt = buildAutoRouterCompressionParams(hydrateAutoRouterCompression(original)); + expect(rebuilt).toEqual(original); + }); +}); diff --git a/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.ts b/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.ts new file mode 100644 index 00000000000..49180d5b4ce --- /dev/null +++ b/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.ts @@ -0,0 +1,52 @@ +/** + * Maps the auto router's compression form state to the two flat litellm_params keys + * the backend reads (litellm.proxy.guardrails.auto_router_compression), and back. + * + * `routing` being undefined means the section was never touched: both keys are + * omitted from the payload, and the request's own compression guardrails apply to + * both hops unchanged. Once `routing` has a value (a guardrail name, or the "none" + * sentinel for explicit no-compression), the auto router is authoritative and the + * model side always gets a concrete value too, mirroring `routing` when same-as + * is chosen and defaulting to "none" otherwise. + */ + +export const NO_COMPRESSION = "none"; + +export interface AutoRouterCompressionState { + routing: string | undefined; + sameAsRouting: boolean; + model: string | undefined; +} + +export interface AutoRouterCompressionLitellmParams { + auto_router_routing_compression?: string; + auto_router_model_compression?: string; +} + +export const DEFAULT_AUTO_ROUTER_COMPRESSION: AutoRouterCompressionState = { + routing: undefined, + sameAsRouting: true, + model: undefined, +}; + +export const buildAutoRouterCompressionParams = ( + state: AutoRouterCompressionState, +): AutoRouterCompressionLitellmParams => { + if (state.routing === undefined) return {}; + return { + auto_router_routing_compression: state.routing, + auto_router_model_compression: state.sameAsRouting ? state.routing : (state.model ?? NO_COMPRESSION), + }; +}; + +export const hydrateAutoRouterCompression = (litellmParams: { + auto_router_routing_compression?: string | null; + auto_router_model_compression?: string | null; +}): AutoRouterCompressionState => { + const routing = litellmParams.auto_router_routing_compression ?? undefined; + if (routing === undefined) return DEFAULT_AUTO_ROUTER_COMPRESSION; + + const model = litellmParams.auto_router_model_compression ?? undefined; + const sameAsRouting = model === undefined || model === routing; + return { routing, sameAsRouting, model: sameAsRouting ? undefined : model }; +}; diff --git a/ui/litellm-dashboard/src/components/add_model/handle_add_auto_router_submit.tsx b/ui/litellm-dashboard/src/components/add_model/handle_add_auto_router_submit.tsx index 9385836ce1a..59d9ecf205e 100644 --- a/ui/litellm-dashboard/src/components/add_model/handle_add_auto_router_submit.tsx +++ b/ui/litellm-dashboard/src/components/add_model/handle_add_auto_router_submit.tsx @@ -1,8 +1,9 @@ import { modelCreateCall } from "../networking"; import { toast } from "@/lib/toast"; import type { ComplexityRouterConfigPayload } from "./build_complexity_router_config"; +import type { AutoRouterCompressionLitellmParams } from "./buildAutoRouterCompression"; -export interface AddAutoRouterValues { +export interface AddAutoRouterValues extends AutoRouterCompressionLitellmParams { auto_router_name: string; auto_router_default_model: string | undefined; model_type: "complexity_router"; @@ -24,6 +25,8 @@ export const handleAddAutoRouterSubmit = async ( model: "auto_router/complexity_router", complexity_router_config: values.complexity_router_config, complexity_router_default_model: values.auto_router_default_model, + auto_router_routing_compression: values.auto_router_routing_compression, + auto_router_model_compression: values.auto_router_model_compression, }, model_info: { ...(values.team_id ? { team_id: values.team_id } : {}), diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.tsx b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.tsx index 970bcaa545f..b93db8d963e 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.tsx +++ b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.tsx @@ -1029,3 +1029,84 @@ describe("EditAutoRouterModal with a stored custom tier set", () => { expect(savedConfig().tier_model_configs).toEqual(CUSTOM_STORED.tier_model_configs); }); }); + +describe("EditAutoRouterModal prompt compression", () => { + beforeEach(() => { + modelPatchUpdateCall.mockClear(); + }); + + const savedLitellmParams = () => { + const [, payload] = modelPatchUpdateCall.mock.calls.at(-1) ?? []; + return payload?.litellm_params; + }; + + const renderWithStoredCompression = ( + compression?: { auto_router_routing_compression?: string; auto_router_model_compression?: string }, + ) => + renderWithProviders( + , + ); + + it("leaves both compression keys out of an untouched save when none were stored", async () => { + const user = userEvent.setup(); + renderWithStoredCompression(); + + await user.click(await screen.findByRole("button", { name: /save changes/i })); + + await waitFor(() => expect(modelPatchUpdateCall).toHaveBeenCalled()); + expect(savedLitellmParams()).not.toHaveProperty("auto_router_routing_compression"); + expect(savedLitellmParams()).not.toHaveProperty("auto_router_model_compression"); + }); + + it("preserves a stored same-as-routing compression through an untouched open-and-save", async () => { + const user = userEvent.setup(); + renderWithStoredCompression({ + auto_router_routing_compression: "headroom-a", + auto_router_model_compression: "headroom-a", + }); + + await user.click(await screen.findByRole("button", { name: /save changes/i })); + + await waitFor(() => expect(modelPatchUpdateCall).toHaveBeenCalled()); + expect(savedLitellmParams()?.auto_router_routing_compression).toBe("headroom-a"); + expect(savedLitellmParams()?.auto_router_model_compression).toBe("headroom-a"); + }); + + it("shows a stored different-compression choice as Use a different compression, not Same", async () => { + const user = userEvent.setup(); + renderWithStoredCompression({ + auto_router_routing_compression: "headroom-a", + auto_router_model_compression: "none", + }); + + await user.click(await screen.findByText("Advanced: Compression")); + + expect(await screen.findByRole("combobox", { name: "Routing decision compression" })).toHaveValue("headroom-a"); + expect(screen.getByRole("radio", { name: "Use a different compression" })).toBeChecked(); + expect(screen.getByRole("combobox", { name: "Model call compression" })).toHaveValue("None (no compression)"); + }); + + it("preserves a stored different-compression choice through an untouched open-and-save", async () => { + const user = userEvent.setup(); + renderWithStoredCompression({ + auto_router_routing_compression: "headroom-a", + auto_router_model_compression: "none", + }); + + await user.click(await screen.findByRole("button", { name: /save changes/i })); + + await waitFor(() => expect(modelPatchUpdateCall).toHaveBeenCalled()); + expect(savedLitellmParams()?.auto_router_routing_compression).toBe("headroom-a"); + expect(savedLitellmParams()?.auto_router_model_compression).toBe("none"); + }); +}); diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx index ea1e5cba6a3..a4852f9e784 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx +++ b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx @@ -41,6 +41,12 @@ import { } from "../add_model/build_complexity_router_config"; import { KeywordTierRule } from "../add_model/KeywordTierRules"; import { DEFAULT_MATCH_THRESHOLD } from "../add_model/SemanticKeywordMatching"; +import { + type AutoRouterCompressionState, + buildAutoRouterCompressionParams, + DEFAULT_AUTO_ROUTER_COMPRESSION, + hydrateAutoRouterCompression, +} from "../add_model/buildAutoRouterCompression"; import { hydrateKeywordTierRules } from "../add_model/complexity_router_keywords"; import { hydrateDimensionWeights, @@ -424,6 +430,9 @@ const EditAutoRouterModal: React.FC = ({ const [semanticMatchingEnabled, setSemanticMatchingEnabled] = useState(false); const [embeddingModel, setEmbeddingModel] = useState(undefined); const [matchThreshold, setMatchThreshold] = useState(DEFAULT_MATCH_THRESHOLD); + const [autoRouterCompression, setAutoRouterCompression] = useState( + DEFAULT_AUTO_ROUTER_COMPRESSION, + ); const [complexityRouterConfig, setComplexityRouterConfig] = useState({ tiers: { SIMPLE: [], MEDIUM: [], COMPLEX: [], REASONING: [] }, classifier_type: "heuristic", @@ -516,6 +525,12 @@ const EditAutoRouterModal: React.FC = ({ setMatchThreshold( typeof parsedConfig.match_threshold === "number" ? parsedConfig.match_threshold : DEFAULT_MATCH_THRESHOLD, ); + setAutoRouterCompression( + hydrateAutoRouterCompression({ + auto_router_routing_compression: modelData.litellm_params?.auto_router_routing_compression, + auto_router_model_compression: modelData.litellm_params?.auto_router_model_compression, + }), + ); form.reset({ ...EMPTY_FORM_VALUES, @@ -628,6 +643,7 @@ const EditAutoRouterModal: React.FC = ({ ...modelData.litellm_params, complexity_router_config: updatedConfig, complexity_router_default_model: defaultModel, + ...buildAutoRouterCompressionParams(autoRouterCompression), }; const updatedModelInfo = { ...modelData.model_info, @@ -749,6 +765,8 @@ const EditAutoRouterModal: React.FC = ({ onMatchThresholdChange={setMatchThreshold} escalationKeywords={escalationKeywords} onEscalationKeywordsChange={setEscalationKeywords} + autoRouterCompression={autoRouterCompression} + onAutoRouterCompressionChange={setAutoRouterCompression} /> ) : ( diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 5f3cca49644..8f6a0700517 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -29222,6 +29222,10 @@ export interface components { auto_router_embedding_model?: string | null; /** Auto Router Max Input Chars */ auto_router_max_input_chars?: number | null; + /** Auto Router Model Compression */ + auto_router_model_compression?: string | null; + /** Auto Router Routing Compression */ + auto_router_routing_compression?: string | null; /** Aws Access Key Id */ aws_access_key_id?: string | null; /** Aws Batch Role Arn */ @@ -39275,6 +39279,10 @@ export interface components { auto_router_embedding_model?: string | null; /** Auto Router Max Input Chars */ auto_router_max_input_chars?: number | null; + /** Auto Router Model Compression */ + auto_router_model_compression?: string | null; + /** Auto Router Routing Compression */ + auto_router_routing_compression?: string | null; /** Aws Access Key Id */ aws_access_key_id?: string | null; /** Aws Batch Role Arn */ From 9e286fe94bf18d29b7bb56e3c2f77d114c10acc5 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Fri, 4 Sep 2026 16:42:11 -0700 Subject: [PATCH 02/13] fix(auto-router): close review findings on per-hop compression - Suppression markers now carry the per-process token `_pre_call_marker` already uses, so a caller cannot switch off an always-on PII, content-filter or compression guardrail by naming it in its own request metadata. - Routing set to "none" with the model side compressed now classifies on the pre-compression snapshot instead of the model-side guardrail's output. - Both the proxy's pre-call arming and the router's routing hook resolve the policy through one tag-aware `policy_for_model`, so an alias with several tag-scoped markers can no longer suppress one marker's guardrail and then route under another marker's policy. - The pre-compression snapshot moved from request metadata to a ContextVar: `refresh_proxy_server_request_body_snapshot` copies metadata into `proxy_server_request.body`, which deployments persist, and the snapshot holds the prompt as it was before any masking guardrail rewrote it. - The compression selector lists Compresr guardrails too, not just Headroom. --- litellm/integrations/custom_guardrail.py | 22 ++- .../guardrails/auto_router_compression.py | 148 +++++++++------ litellm/router.py | 32 ++-- .../integrations/test_custom_guardrail.py | 35 +++- .../test_auto_router_compression.py | 177 +++++++++++++----- tests/test_litellm/test_router.py | 29 +++ .../add_model/CompressionControls.tsx | 5 +- .../add_model/buildAutoRouterCompression.ts | 7 + 8 files changed, 322 insertions(+), 133 deletions(-) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 7f8effa2317..558e97cfc16 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -941,17 +941,29 @@ class CustomGuardrail(CustomLogger): """ return False - def _suppressed_by_auto_router_compression(self, data: dict) -> bool: - """True when an auto router's own compression policy suppresses this guardrail. + def auto_router_suppression_marker(self) -> str | None: + """The value `arm_pre_call` must write to suppress this guardrail. - Set only by litellm.proxy.guardrails.auto_router_compression.arm_pre_call, never - by the caller, so a request cannot suppress its own guardrails this way. + Carries the per-process token for the same reason `_pre_call_marker` does: a + caller controls request metadata, so a bare guardrail name there would let any + request switch off a PII, content-filter, or compression guardrail for itself. + The token is never sent to the caller, so the marker cannot be forged. """ + name: Final = self.guardrail_name + if not name: + return None + return f"{_PRE_CALL_EXECUTED_TOKEN}:{name}" + + def _suppressed_by_auto_router_compression(self, data: dict[str, object]) -> bool: + """True when an auto router's own compression policy suppresses this guardrail.""" + marker: Final = self.auto_router_suppression_marker() + if marker is None: + return False for meta_key in ("metadata", "litellm_metadata"): meta = data.get(meta_key) if isinstance(meta, dict): suppressed = meta.get(AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY) - if isinstance(suppressed, list) and self.guardrail_name in suppressed: + if isinstance(suppressed, list) and marker in suppressed: return True return False diff --git a/litellm/proxy/guardrails/auto_router_compression.py b/litellm/proxy/guardrails/auto_router_compression.py index c3fba937d22..7ccd1937543 100644 --- a/litellm/proxy/guardrails/auto_router_compression.py +++ b/litellm/proxy/guardrails/auto_router_compression.py @@ -12,34 +12,32 @@ guardrail is suppressed for that request, and only these two settings decide wha each hop sees. """ -import copy -from collections.abc import Mapping +import contextvars +from collections.abc import Mapping, MutableMapping, Sequence from dataclasses import dataclass from typing import TYPE_CHECKING, Any, Final from litellm._logging import verbose_proxy_logger from litellm.constants import AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY -from litellm.litellm_core_utils.core_helpers import ( - get_metadata_variable_name_from_kwargs, - get_or_create_metadata_bucket, -) +from litellm.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 -else: - CustomGuardrail = Any - Router = Any COMPRESSION_GUARDRAIL_PROVIDERS: Final = frozenset({"headroom", "compresr"}) _NO_COMPRESSION: Final = "none" -# Metadata key stashing the pre-compression messages so a routing decision that -# names a different compression than the model call still compresses the -# original text, not whatever the model-side guardrail already rewrote it to. -AUTO_ROUTER_ROUTING_MESSAGES_SNAPSHOT_KEY: Final = "_auto_router_routing_messages_snapshot" +# The pre-compression messages, so a routing decision that does not share the model +# call's compression still classifies on the original text. Deliberately a ContextVar +# rather than a metadata key: `refresh_proxy_server_request_body_snapshot` copies +# metadata into `proxy_server_request.body`, which deployments persist to spend logs, +# and this holds the prompt as it was before any masking guardrail rewrote it. +_routing_messages_snapshot: Final[contextvars.ContextVar[tuple[Mapping[str, object], ...] | None]] = ( + contextvars.ContextVar("litellm_auto_router_routing_messages_snapshot", default=None) +) @dataclass(frozen=True, slots=True) @@ -72,30 +70,51 @@ def policy_from_litellm_params(litellm_params: Mapping[str, object]) -> AutoRout def policy_for_model( - llm_router: "Router | None", model_alias: str, team_id: str | None + llm_router: "Router | None", + model_alias: str, + team_id: str | None, + request_tags: Sequence[str], ) -> AutoRouterCompressionPolicy | None: - """The compression policy declared by the auto router marker deployment `model_alias` resolves to. + """The compression policy of the auto router marker `model_alias` resolves to. - Mirrors the alias lookup in ``_check_and_merge_model_level_guardrails``: this runs - before routing has picked a strategy, so it takes the first marker deployment for - the alias rather than disambiguating by request tags. + Both the proxy's pre-call arming and the router's routing hook resolve the policy + through here, with the same tag rule, so an alias carrying several tag-scoped + markers can never suppress one marker's guardrail and then route under another + marker's policy. """ if llm_router is None: return None deployments: Final = llm_router.get_model_list(model_name=model_alias, team_id=team_id) or [] - for deployment in deployments: - litellm_params: Final = deployment.get("litellm_params") or {} - model_field = litellm_params.get("model") - if not isinstance(model_field, str) or not model_field.startswith(AUTO_ROUTER_MODEL_PREFIX): - continue - policy = policy_from_litellm_params(litellm_params) + markers: Final = tuple( + litellm_params + for deployment in deployments + if isinstance(litellm_params := deployment.get("litellm_params") or {}, 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 requested.issuperset(frozenset(params.get("tags") or ())) + ) + for params in (*tag_matched, *markers): + policy = policy_from_litellm_params(params) if policy is not None: return policy return None -def _active_compression_guardrail_names() -> frozenset[str]: - """Names of every currently-active guardrail whose type is a compression guardrail.""" +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 _active_compression_guardrails() -> tuple["CustomGuardrail", ...]: + """Every currently-active guardrail whose type is a compression guardrail.""" import litellm from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.proxy.guardrails.guardrail_registry import guardrail_class_registry @@ -104,21 +123,20 @@ def _active_compression_guardrail_names() -> frozenset[str]: cls for name, cls in guardrail_class_registry.items() if name in COMPRESSION_GUARDRAIL_PROVIDERS ) if not compression_classes: - return frozenset() + return () active: Final = litellm.logging_callback_manager.get_custom_loggers_for_type(callback_type=CustomGuardrail) - return frozenset( - cb.guardrail_name for cb in active if isinstance(cb, compression_classes) and cb.guardrail_name - ) + return tuple(cb for cb in active if isinstance(cb, compression_classes) and cb.guardrail_name) -async def arm_pre_call(data: dict, llm_router: "Router | None") -> dict: +async def arm_pre_call(data: MutableMapping[str, object], llm_router: "Router | None") -> MutableMapping[str, object]: """Apply an auto router's compression policy, if any, before guardrails run. Suppresses every other compression guardrail, re-enables the model-side guardrail the policy names (if any) even when it isn't ``default_on``, and - snapshots the pre-compression messages so the routing decision can compress - them independently of whatever the model-side guardrail does to `data`. + snapshots the pre-compression messages so the routing decision can read them + independently of whatever the model-side guardrail does to `data`. """ + _routing_messages_snapshot.set(None) if llm_router is None: return data @@ -129,21 +147,27 @@ async def arm_pre_call(data: dict, llm_router: "Router | None") -> dict: # Read-only until a policy is confirmed: creating the metadata bucket for every # request, including the vast majority with no auto-router compression policy, # would be an unwanted side effect of merely checking for one. - metadata_key: Final = get_metadata_variable_name_from_kwargs(data) - existing_bucket: Final = data.get(metadata_key) - other_bucket: Final = data.get("metadata" if metadata_key == "litellm_metadata" else "litellm_metadata") - team_id: Final = (existing_bucket.get("user_api_key_team_id") if isinstance(existing_bucket, dict) else None) or ( - other_bucket.get("user_api_key_team_id") if isinstance(other_bucket, dict) else None - ) + from litellm.router_strategy.tag_based_routing import _get_tags_from_request_kwargs - policy: Final = policy_for_model(llm_router=llm_router, model_alias=model_alias, team_id=team_id) + 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 data _, metadata = get_or_create_metadata_bucket(data) - suppressed: Final = _active_compression_guardrail_names() - ({policy.model} if policy.model else set()) + # Markers carry a per-process token so a caller cannot suppress a guardrail by + # naming it in its own request metadata. + suppressed: Final = tuple( + marker + for guardrail in _active_compression_guardrails() + if guardrail.guardrail_name != policy.model and (marker := guardrail.auto_router_suppression_marker()) + ) if suppressed: - metadata[AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY] = sorted(suppressed) + metadata[AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY] = list(suppressed) if policy.model is not None: requested = metadata.get("guardrails") @@ -157,44 +181,52 @@ async def arm_pre_call(data: dict, llm_router: "Router | None") -> dict: snapshot: Final = resolve_structured_messages(messages=data.get("messages"), request_kwargs=data) if snapshot is not None: - metadata[AUTO_ROUTER_ROUTING_MESSAGES_SNAPSHOT_KEY] = copy.deepcopy(snapshot) + _routing_messages_snapshot.set(tuple(dict(message) for message in snapshot)) return data +def _snapshot_messages() -> list[dict[str, Any]] | None: + snapshot: Final = _routing_messages_snapshot.get() + return None if snapshot is None else [dict(message) for message in snapshot] + + async def messages_for_routing( policy: AutoRouterCompressionPolicy | None, messages: list[dict[str, Any]] | None, request_kwargs: Mapping[str, object], ) -> list[dict[str, Any]] | None: - """Messages to use for a routing decision, compressed per `policy.routing`. + """Messages to use for a routing decision, per `policy.routing`. - Returns None when there is no policy or the policy's routing side names no - compression, meaning the caller should route on whatever messages it already - has. The model call is untouched by this function either way: model-side - compression, if any, already ran as an ordinary pre-call guardrail before the - router was ever reached. + Returns None when the caller should route on whatever messages it already has. + The model call is untouched either way: model-side compression, if any, already + ran as an ordinary pre-call guardrail before the router was reached, so when the + two hops differ the routing decision reads the pre-compression snapshot rather + than what that guardrail left behind. """ - if policy is None or policy.routing is None: + if policy is None: + return None + + original: Final = _snapshot_messages() or messages + + if policy.routing is None: + # Explicitly no compression for routing. When the model side compressed, the + # messages in hand are its output, so fall back to the untouched snapshot. + return _snapshot_messages() if policy.model is not None else None + + if not original: return None from litellm.proxy.common_utils.registry_read_through import ( get_initialized_guardrail_with_read_through, ) - metadata_key: Final = get_metadata_variable_name_from_kwargs(request_kwargs) - metadata: Final = request_kwargs.get(metadata_key) - snapshot: Final = metadata.get(AUTO_ROUTER_ROUTING_MESSAGES_SNAPSHOT_KEY) if isinstance(metadata, dict) else None - original: Final = snapshot if isinstance(snapshot, list) else messages - if not original: - return None - guardrail: Final = await get_initialized_guardrail_with_read_through(policy.routing) if guardrail is None: verbose_proxy_logger.warning( "AutoRouter compression: guardrail '%s' not found; routing on uncompressed messages", policy.routing ) - return None + return original inputs: GenericGuardrailAPIInputs = {"structured_messages": [dict(m) for m in original]} # A throwaway request_data: apply_guardrail writes its stats onto this dict, not diff --git a/litellm/router.py b/litellm/router.py index 8149ec60ddc..bcb2e2aa7ff 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -13039,11 +13039,19 @@ class Router: from litellm.proxy.guardrails.auto_router_compression import ( messages_for_routing, - policy_from_litellm_params, + policy_for_model, + team_id_from_request, ) - marker_params: Final = self._alias_marker_litellm_params(registered_model_name, selected_strategy.tags) - compression_policy: Final = policy_from_litellm_params(marker_params) if marker_params else None + # Resolved through the same tag-aware lookup the proxy's pre-call arming used, + # so an alias carrying several tag-scoped markers cannot suppress one marker's + # guardrail and then route under a different marker's policy. + 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), + ) # When both hops share the same compression, the model-side guardrail already # ran in the proxy's ordinary pre-call hook and compressed `messages` in place # (arm_pre_call armed it whether or not it is `default_on`); reuse that result @@ -13133,16 +13141,9 @@ class Router: return pre_routing_hook_response - def _alias_marker_litellm_params( + def _forwardable_alias_marker_params( self, model: str, strategy_tags: tuple[str, ...] - ) -> Mapping[str, object] | None: - """The auto-router marker deployment's own `litellm_params` for `model`, tag-scoped. - - Shared by `_forwardable_alias_marker_params` (forwarding api_base/api_key/... - gaps onto the routed deployment) and the auto-router compression policy lookup - (reading `auto_router_routing_compression`/`auto_router_model_compression`), so - both read the same marker row when an alias has more than one, tag-scoped marker. - """ + ) -> tuple[tuple[str, object], ...]: marker_params: Final = tuple( litellm_params for idx in self.model_name_to_deployment_indices.get(model, ()) @@ -13152,12 +13153,7 @@ class Router: tag_matched: Final = tuple( params for params in marker_params if tuple(params.get("tags") or ()) == strategy_tags ) - return tag_matched[0] if tag_matched else (marker_params[0] if marker_params else None) - - def _forwardable_alias_marker_params( - self, model: str, strategy_tags: tuple[str, ...] - ) -> tuple[tuple[str, object], ...]: - selected: Final = self._alias_marker_litellm_params(model, strategy_tags) + selected: Final = tag_matched[0] if tag_matched else (marker_params[0] if marker_params else None) if selected is None: return () return tuple( diff --git a/tests/test_litellm/integrations/test_custom_guardrail.py b/tests/test_litellm/integrations/test_custom_guardrail.py index 1fb4299cb56..f590903cb74 100644 --- a/tests/test_litellm/integrations/test_custom_guardrail.py +++ b/tests/test_litellm/integrations/test_custom_guardrail.py @@ -532,7 +532,9 @@ class TestCustomGuardrailShouldRunGuardrail: data = { "model": "smart-router", "metadata": { - AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY: ["headroom-default"], + AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY: [ + always_on.auto_router_suppression_marker() + ], }, } @@ -550,10 +552,13 @@ class TestCustomGuardrailShouldRunGuardrail: default_on=True, event_hook=GuardrailEventHooks.pre_call, ) + other = CustomGuardrail(guardrail_name="some-other-guardrail", default_on=True) data = { "model": "smart-router", "metadata": { - AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY: ["some-other-guardrail"], + AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY: [ + other.auto_router_suppression_marker() + ], }, } @@ -562,6 +567,32 @@ class TestCustomGuardrailShouldRunGuardrail: is True ) + def test_should_run_guardrail_ignores_a_forged_suppression_marker(self): + """A caller controls request metadata, so a bare guardrail name there must not + switch off an always-on guardrail: only the per-process marker counts.""" + from litellm.constants import AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY + from litellm.types.guardrails import GuardrailEventHooks + + always_on = CustomGuardrail( + guardrail_name="headroom-default", + default_on=True, + event_hook=GuardrailEventHooks.pre_call, + ) + forged = { + "model": "smart-router", + "metadata": { + AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY: [ + "headroom-default", + "forged-token:headroom-default", + ], + }, + } + + assert ( + always_on.should_run_guardrail(data=forged, event_type=GuardrailEventHooks.pre_call) + is True + ) + class TestApplyGuardrailCheck: def test_apply_guardrail_check_only_on_direct_implementation(self): diff --git a/tests/test_litellm/proxy/guardrails/test_auto_router_compression.py b/tests/test_litellm/proxy/guardrails/test_auto_router_compression.py index b2e83e75768..b906e60bb86 100644 --- a/tests/test_litellm/proxy/guardrails/test_auto_router_compression.py +++ b/tests/test_litellm/proxy/guardrails/test_auto_router_compression.py @@ -4,28 +4,33 @@ Unit tests for litellm.proxy.guardrails.auto_router_compression. Covers: - policy_from_litellm_params: absent keys mean no policy; the "none" sentinel normalizes to explicit no-compression within an active policy; is_same -- policy_for_model: finds the auto-router marker deployment for an alias -- arm_pre_call: no-op without a policy; suppresses active compression guardrails; - arms the model-side guardrail even when it isn't default_on; snapshots messages -- messages_for_routing: no-op without a policy or an unset routing side; compresses - via the named guardrail's apply_guardrail; never writes stats onto the caller's - own request_kwargs (regression for double-counted compression savings) +- policy_for_model: finds the auto-router marker deployment for an alias, and + picks the tag-scoped marker the request's tags actually match +- arm_pre_call: no-op without a policy; suppresses active compression guardrails + with a forgery-proof marker; arms the model-side guardrail even when it isn't + default_on; keeps the pre-compression snapshot out of persisted metadata +- messages_for_routing: no-op without a policy; routes on the pre-compression + snapshot when the two hops differ; compresses via the named guardrail's + apply_guardrail; never writes stats onto the caller's own request_kwargs + (regression for double-counted compression savings) """ +import json from typing import Any import pytest +from litellm.constants import AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.proxy.guardrails import auto_router_compression from litellm.proxy.guardrails.auto_router_compression import ( - AUTO_ROUTER_ROUTING_MESSAGES_SNAPSHOT_KEY, AutoRouterCompressionPolicy, arm_pre_call, messages_for_routing, policy_for_model, policy_from_litellm_params, ) -from litellm.constants import AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY +from litellm.types.guardrails import GuardrailEventHooks from litellm.types.utils import GenericGuardrailAPIInputs @@ -76,36 +81,61 @@ class _FakeRouter: 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) is None + 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) is None + 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) is None + 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( [ - { - "model_name": "smart-router", - "litellm_params": { - "model": "auto_router/complexity_router", - "auto_router_routing_compression": "headroom-a", - "auto_router_model_compression": "none", - }, - } + _marker({"auto_router_routing_compression": "headroom-eu"}, tags=["eu"]), + _marker({"auto_router_routing_compression": "headroom-us"}, tags=["us"]), ] ) - policy = policy_for_model(llm_router=router, model_alias="smart-router", team_id=None) + + 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) @@ -186,10 +216,25 @@ class TestArmPreCall: data = {"model": "smart-router", "messages": [{"role": "user", "content": "hi"}]} result = await arm_pre_call(data=data, llm_router=router) suppressed = result["metadata"][AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY] - assert "always-on-compression" in suppressed + assert suppressed == [always_on.auto_router_suppression_marker()] + # The bare name alone must never suppress: that is what a caller could forge. + assert "always-on-compression" not in suppressed + assert always_on.should_run_guardrail(data=result, 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_a_caller_cannot_suppress_a_guardrail_by_naming_it_in_metadata(self): + """Regression: request metadata is caller-controlled, so a bare guardrail name + there must not switch off a PII, content-filter, or compression guardrail.""" + guardrail = _RecordingCompressionGuardrail(guardrail_name="always-on-compression") + forged = { + "model": "smart-router", + "metadata": {AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY: ["always-on-compression"]}, + } + + assert guardrail._suppressed_by_auto_router_compression(forged) is False + @pytest.mark.asyncio async def test_model_side_guardrail_is_requested_even_when_not_default_on(self): router = _FakeRouter( @@ -209,43 +254,82 @@ class TestArmPreCall: assert result["metadata"]["guardrails"] == ["headroom-b"] @pytest.mark.asyncio - async def test_snapshots_original_messages(self): - router = _FakeRouter( - [ - { - "model_name": "smart-router", - "litellm_params": { - "model": "auto_router/complexity_router", - "auto_router_routing_compression": "headroom-a", - "auto_router_model_compression": "none", - }, - } - ] - ) - original_messages = [{"role": "user", "content": "hi"}] + async def test_snapshot_never_lands_in_persisted_metadata(self): + """Regression: refresh_proxy_server_request_body_snapshot copies metadata into + proxy_server_request.body, which deployments persist to spend logs. The + pre-compression snapshot holds the prompt before any masking guardrail ran, so + it must live outside anything that gets serialized.""" + router = _FakeRouter([_marker({"auto_router_routing_compression": "headroom-a"})]) + original_messages = [{"role": "user", "content": "my ssn is 123-45-6789"}] data = {"model": "smart-router", "messages": original_messages} + result = await arm_pre_call(data=data, llm_router=router) - snapshot = result["metadata"][AUTO_ROUTER_ROUTING_MESSAGES_SNAPSHOT_KEY] - assert snapshot == original_messages - assert snapshot is not original_messages # a copy, not the live reference + + assert "123-45-6789" not in json.dumps(result["metadata"]) + assert auto_router_compression._snapshot_messages() == original_messages + + @pytest.mark.asyncio + async def test_snapshot_is_a_copy_not_the_live_message_list(self): + router = _FakeRouter([_marker({"auto_router_routing_compression": "headroom-a"})]) + original_messages = [{"role": "user", "content": "hi"}] + + await arm_pre_call(data={"model": "smart-router", "messages": original_messages}, llm_router=router) + original_messages[0]["content"] = "mutated after the snapshot" + + assert auto_router_compression._snapshot_messages() == [{"role": "user", "content": "hi"}] + + @pytest.mark.asyncio + async def test_a_request_without_a_policy_clears_a_previous_snapshot(self): + router_with = _FakeRouter([_marker({"auto_router_routing_compression": "headroom-a"})]) + await arm_pre_call(data={"model": "smart-router", "messages": [{"role": "user", "content": "first"}]}, + llm_router=router_with) + + router_without = _FakeRouter( + [{"model_name": "plain", "litellm_params": {"model": "openai/gpt-4o-mini"}}] + ) + await arm_pre_call(data={"model": "plain", "messages": [{"role": "user", "content": "second"}]}, + llm_router=router_without) + + assert auto_router_compression._snapshot_messages() is None class TestMessagesForRouting: + @pytest.fixture(autouse=True) + def _clear_snapshot(self): + auto_router_compression._routing_messages_snapshot.set(None) + yield + auto_router_compression._routing_messages_snapshot.set(None) + @pytest.mark.asyncio async def test_no_policy_returns_none(self): assert await messages_for_routing(policy=None, messages=[], request_kwargs={}) is None @pytest.mark.asyncio - async def test_routing_side_unset_returns_none(self): - policy = AutoRouterCompressionPolicy(routing=None, model="headroom-a") + 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_unknown_guardrail_name_returns_none(self): + async def test_routing_none_with_model_compression_routes_on_the_snapshot(self): + """Regression: with routing explicitly off and the model side compressed, the + messages in hand are the model-side guardrail's output. Routing asked for no + compression, so it must read the pre-compression snapshot instead.""" + original = [{"role": "user", "content": "the full original conversation"}] + auto_router_compression._routing_messages_snapshot.set(tuple(dict(m) for m in original)) + policy = AutoRouterCompressionPolicy(routing=None, model="headroom-a") + model_compressed = [{"role": "user", "content": "[COMPRESSED] the full original conversation"}] + + result = await messages_for_routing(policy=policy, messages=model_compressed, request_kwargs={}) + + assert result == original + + @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 is None + assert result == messages @pytest.mark.asyncio async def test_compresses_via_the_named_guardrail(self, registered_guardrail): @@ -256,15 +340,14 @@ class TestMessagesForRouting: @pytest.mark.asyncio async def test_uses_the_snapshot_when_present(self, registered_guardrail): - policy = AutoRouterCompressionPolicy(routing="fake-compress", model=None) - snapshot = [{"role": "user", "content": "original"}] - request_kwargs = {"metadata": {AUTO_ROUTER_ROUTING_MESSAGES_SNAPSHOT_KEY: snapshot}} - # `messages` here stands in for whatever a model-side guardrail already + policy = AutoRouterCompressionPolicy(routing="fake-compress", model="headroom-b") + auto_router_compression._routing_messages_snapshot.set(({"role": "user", "content": "original"},)) + # `messages` here stands in for whatever the model-side guardrail already # rewrote `data["messages"]` to -- routing must ignore it and compress the # pristine snapshot instead. already_rewritten = [{"role": "user", "content": "rewritten by another guardrail"}] result = await messages_for_routing( - policy=policy, messages=already_rewritten, request_kwargs=request_kwargs + policy=policy, messages=already_rewritten, request_kwargs={} ) assert result == [{"role": "user", "content": "[COMPRESSED] original"}] diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index cdae3b131ae..4c5813911ac 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -10105,6 +10105,35 @@ class TestAutoRouterCompressionDecoupling: assert response.messages == original_messages assert registered_guardrail.call_count == 0 + @pytest.mark.asyncio + async def test_routing_none_routes_on_the_original_not_the_model_compressed_messages( + self, registered_guardrail + ): + """Regression: with routing explicitly off and the model side compressed, the + messages the router holds are the model-side guardrail's output. Routing asked + for no compression, so it has to classify on the pre-compression snapshot.""" + from litellm.proxy.guardrails import auto_router_compression + + router, strategy = self._router( + { + "auto_router_routing_compression": "none", + "auto_router_model_compression": "fake-compress", + } + ) + original_messages = self._messages() + auto_router_compression._routing_messages_snapshot.set(tuple(dict(m) for m in original_messages)) + model_compressed = [{"role": "user", "content": "[COMPRESSED] What is the capital of France?"}] + + try: + await router.async_pre_routing_hook( + model="smart-router", request_kwargs={"metadata": {}}, messages=model_compressed + ) + finally: + auto_router_compression._routing_messages_snapshot.set(None) + + assert strategy.received_messages == original_messages + assert registered_guardrail.call_count == 0 + @pytest.mark.asyncio async def test_same_compression_on_both_hops_compresses_once(self, registered_guardrail): """The same/different distinction exists so a shared choice does not pay for diff --git a/ui/litellm-dashboard/src/components/add_model/CompressionControls.tsx b/ui/litellm-dashboard/src/components/add_model/CompressionControls.tsx index a0a240f76b0..52ea7645034 100644 --- a/ui/litellm-dashboard/src/components/add_model/CompressionControls.tsx +++ b/ui/litellm-dashboard/src/components/add_model/CompressionControls.tsx @@ -5,8 +5,7 @@ import { RadioGroup, RadioGroupItem } from "@/components/ui/radio-group"; import { Info } from "lucide-react"; import React from "react"; import { useGuardrails } from "@/app/(dashboard)/hooks/guardrails/useGuardrails"; -import { COMPRESSION_GUARDRAIL_PROVIDER } from "@/app/(dashboard)/cost-optimization/_components/helpers"; -import { NO_COMPRESSION } from "./buildAutoRouterCompression"; +import { isCompressionGuardrailProvider, NO_COMPRESSION } from "./buildAutoRouterCompression"; interface CompressionControlsProps { routing: string | undefined; @@ -29,7 +28,7 @@ const CompressionControls: React.FC = ({ }) => { const { data } = useGuardrails(); const compressionOptions: SearchSelectOption[] = (data?.guardrails ?? []) - .filter((g) => (g.litellm_params?.guardrail ?? "").toString().toLowerCase() === COMPRESSION_GUARDRAIL_PROVIDER) + .filter((g) => isCompressionGuardrailProvider(g.litellm_params?.guardrail)) .map((g) => ({ label: g.guardrail_name, value: g.guardrail_name })); const options: SearchSelectOption[] = [NONE_OPTION, ...compressionOptions]; diff --git a/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.ts b/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.ts index 49180d5b4ce..5afdcf2b15c 100644 --- a/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.ts +++ b/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.ts @@ -12,6 +12,13 @@ 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; From d0a80067377cb6978deac35492b6681a65c991fc Mon Sep 17 00:00:00 2001 From: moe-berri Date: Fri, 4 Sep 2026 17:03:01 -0700 Subject: [PATCH 03/13] fix(auto-router compression): tag-scoped markers now take precedence over untagged An untagged marker (no tags key or empty tags list) was matching every request because requested.issuperset(frozenset()) is always true. When an alias carried multiple markers, the loop tried tag-matched markers first, but an untagged one could still match the tag-match query, and then the first one with a policy would be returned. Now only markers with a non-empty tags list can match via the tag-specific lookup; untagged markers are tried only after all tag-specific ones. Regression test added: test_tag_scoped_marker_takes_precedence_over_untagged fails with the old code. Also removed unused Any import per greptile's typing note. --- .../proxy/guardrails/auto_router_compression.py | 16 ++++++++++------ .../guardrails/test_auto_router_compression.py | 12 ++++++++++++ 2 files changed, 22 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/guardrails/auto_router_compression.py b/litellm/proxy/guardrails/auto_router_compression.py index 7ccd1937543..490ce550003 100644 --- a/litellm/proxy/guardrails/auto_router_compression.py +++ b/litellm/proxy/guardrails/auto_router_compression.py @@ -15,7 +15,7 @@ each hop sees. import contextvars from collections.abc import Mapping, MutableMapping, Sequence from dataclasses import dataclass -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Final from litellm._logging import verbose_proxy_logger from litellm.constants import AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY @@ -93,7 +93,9 @@ def policy_for_model( ) requested: Final = frozenset(request_tags) tag_matched: Final = tuple( - params for params in markers if requested.issuperset(frozenset(params.get("tags") or ())) + params + for params in markers + if (tags := params.get("tags")) and requested.issuperset(frozenset(tags)) ) for params in (*tag_matched, *markers): policy = policy_from_litellm_params(params) @@ -186,16 +188,16 @@ async def arm_pre_call(data: MutableMapping[str, object], llm_router: "Router | return data -def _snapshot_messages() -> list[dict[str, Any]] | None: +def _snapshot_messages() -> list[dict[str, object]] | None: snapshot: Final = _routing_messages_snapshot.get() return None if snapshot is None else [dict(message) for message in snapshot] async def messages_for_routing( policy: AutoRouterCompressionPolicy | None, - messages: list[dict[str, Any]] | None, + messages: list[dict[str, object]] | None, request_kwargs: Mapping[str, object], -) -> list[dict[str, Any]] | None: +) -> list[dict[str, object]] | None: """Messages to use for a routing decision, per `policy.routing`. Returns None when the caller should route on whatever messages it already has. @@ -228,7 +230,9 @@ async def messages_for_routing( ) return original - inputs: GenericGuardrailAPIInputs = {"structured_messages": [dict(m) for m in original]} + inputs: GenericGuardrailAPIInputs = { + "structured_messages": [dict(m) for m in original] # pyright: ignore[reportAssignmentType] # plain dicts, not AllMessageValues; see headroom.py's own use of this shape + } # A throwaway request_data: apply_guardrail writes its stats onto this dict, not # the real request's metadata, so routing-side compression never double-counts # against extract_compression_saved_tokens's model-savings accounting. diff --git a/tests/test_litellm/proxy/guardrails/test_auto_router_compression.py b/tests/test_litellm/proxy/guardrails/test_auto_router_compression.py index b906e60bb86..79f069d8a2e 100644 --- a/tests/test_litellm/proxy/guardrails/test_auto_router_compression.py +++ b/tests/test_litellm/proxy/guardrails/test_auto_router_compression.py @@ -138,6 +138,18 @@ class TestPolicyForModel: ) assert policy == AutoRouterCompressionPolicy(routing="headroom-a", model=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.""" From e273cf301fc710a88bb3820e3d35f7fe6a6b20bb Mon Sep 17 00:00:00 2001 From: moe-berri Date: Fri, 4 Sep 2026 17:21:04 -0700 Subject: [PATCH 04/13] fix(ci): satisfy ruff format, prettier, and eslint max-lines gates - ruff format on auto_router_compression.py (a long comprehension wrapped across three lines instead of one) - prettier on buildAutoRouterCompression.ts and the two test files it touched - ComplexityRouterConfig.tsx crossed the 800-line eslint max-lines ceiling once the compression accordion entry landed. Extracted TierRowSelect into its own file (already self-contained, used only within this file and PlanModeOverrideControls) and simplified CompressionControls' props to a single state/onChange pair instead of six individual callbacks, moving the per-field derivation into the component that already owns this state shape --- .../guardrails/auto_router_compression.py | 4 +- .../add_model/ComplexityRouterConfig.tsx | 40 +------------------ .../add_model/CompressionControls.tsx | 29 +++++++------- .../components/add_model/TierRowSelect.tsx | 25 ++++++++++++ .../add_model/buildAutoRouterCompression.ts | 2 +- .../edit_auto_router_modal.test.tsx | 7 ++-- 6 files changed, 47 insertions(+), 60 deletions(-) create mode 100644 ui/litellm-dashboard/src/components/add_model/TierRowSelect.tsx diff --git a/litellm/proxy/guardrails/auto_router_compression.py b/litellm/proxy/guardrails/auto_router_compression.py index 490ce550003..2122d353def 100644 --- a/litellm/proxy/guardrails/auto_router_compression.py +++ b/litellm/proxy/guardrails/auto_router_compression.py @@ -93,9 +93,7 @@ def policy_for_model( ) requested: Final = frozenset(request_tags) tag_matched: Final = tuple( - params - for params in markers - if (tags := params.get("tags")) and requested.issuperset(frozenset(tags)) + params for params in markers if (tags := params.get("tags")) and requested.issuperset(frozenset(tags)) ) for params in (*tag_matched, *markers): policy = policy_from_litellm_params(params) diff --git a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx index 42115265034..92fd9893995 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"; @@ -370,27 +370,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>; @@ -889,22 +868,7 @@ const ComplexityRouterConfig: React.FC = ({ key: "compression", label: Advanced: Compression, children: ( - - onAutoRouterCompressionChange({ - ...autoRouterCompression, - routing, - sameAsRouting: routing === undefined ? true : autoRouterCompression.sameAsRouting, - }) - } - sameAsRouting={autoRouterCompression.sameAsRouting} - onSameAsRoutingChange={(sameAsRouting) => - onAutoRouterCompressionChange({ ...autoRouterCompression, sameAsRouting }) - } - model={autoRouterCompression.model} - onModelChange={(model) => onAutoRouterCompressionChange({ ...autoRouterCompression, model })} - /> + ), }, ] diff --git a/ui/litellm-dashboard/src/components/add_model/CompressionControls.tsx b/ui/litellm-dashboard/src/components/add_model/CompressionControls.tsx index 52ea7645034..c1817918f60 100644 --- a/ui/litellm-dashboard/src/components/add_model/CompressionControls.tsx +++ b/ui/litellm-dashboard/src/components/add_model/CompressionControls.tsx @@ -5,27 +5,26 @@ 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 { isCompressionGuardrailProvider, NO_COMPRESSION } from "./buildAutoRouterCompression"; +import { + AutoRouterCompressionState, + isCompressionGuardrailProvider, + NO_COMPRESSION, +} from "./buildAutoRouterCompression"; interface CompressionControlsProps { - routing: string | undefined; - onRoutingChange: (value: string | undefined) => void; - sameAsRouting: boolean; - onSameAsRoutingChange: (same: boolean) => void; - model: string | undefined; - onModelChange: (value: string | undefined) => void; + value: AutoRouterCompressionState; + onChange: (state: AutoRouterCompressionState) => void; } const NONE_OPTION: SearchSelectOption = { label: "None (no compression)", value: NO_COMPRESSION }; -const CompressionControls: React.FC = ({ - routing, - onRoutingChange, - sameAsRouting, - onSameAsRoutingChange, - model, - onModelChange, -}) => { +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)) 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/buildAutoRouterCompression.ts b/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.ts index 5afdcf2b15c..c86416b507f 100644 --- a/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.ts +++ b/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.ts @@ -42,7 +42,7 @@ export const buildAutoRouterCompressionParams = ( if (state.routing === undefined) return {}; return { auto_router_routing_compression: state.routing, - auto_router_model_compression: state.sameAsRouting ? state.routing : (state.model ?? NO_COMPRESSION), + auto_router_model_compression: state.sameAsRouting ? state.routing : state.model ?? NO_COMPRESSION, }; }; 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 b93db8d963e..d8474b492e1 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 @@ -1040,9 +1040,10 @@ describe("EditAutoRouterModal prompt compression", () => { return payload?.litellm_params; }; - const renderWithStoredCompression = ( - compression?: { auto_router_routing_compression?: string; auto_router_model_compression?: string }, - ) => + const renderWithStoredCompression = (compression?: { + auto_router_routing_compression?: string; + auto_router_model_compression?: string; + }) => renderWithProviders( Date: Fri, 4 Sep 2026 17:42:52 -0700 Subject: [PATCH 05/13] refactor(auto-router compression): satisfy the LIT001/LIT002 type-discipline gate The gate has no headroom, so the new module had to stop introducing mutable collections rather than spend budget on them: - the marker lookup falls back to () and drops an `or {}` that isinstance already covered - the suppression list is stored as the tuple it was built as; the read side in custom_guardrail accepts list or tuple, since JSON round-trips it to a list - the snapshot holds MappingProxyType entries, so it is immutable at rest and _snapshot_messages can hand back the stored tuple with no defensive copy - arm_pre_call returns None instead of echoing back the dict it mutates in place - _suppressed_by_auto_router_compression takes a Mapping, which is all it reads The four remaining mutable spots are external contracts, each suppressed with the reason: the pre-routing hook protocol types messages as list[dict], the metadata["guardrails"] key is extended by litellm_pre_call_utils via an isinstance(..., list) check, apply_guardrail takes a dict it writes stats into, and pydantic's model_copy takes a dict. --- litellm/integrations/custom_guardrail.py | 8 +- litellm/proxy/common_request_processing.py | 2 +- .../guardrails/auto_router_compression.py | 85 +++++++++++-------- litellm/router.py | 3 +- .../test_auto_router_compression.py | 65 ++++++-------- 5 files changed, 83 insertions(+), 80 deletions(-) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 558e97cfc16..1ec08641706 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -954,16 +954,18 @@ class CustomGuardrail(CustomLogger): return None return f"{_PRE_CALL_EXECUTED_TOKEN}:{name}" - def _suppressed_by_auto_router_compression(self, data: dict[str, object]) -> bool: + def _suppressed_by_auto_router_compression(self, data: Mapping[str, object]) -> bool: """True when an auto router's own compression policy suppresses this guardrail.""" marker: Final = self.auto_router_suppression_marker() if marker is None: return False for meta_key in ("metadata", "litellm_metadata"): meta = data.get(meta_key) - if isinstance(meta, dict): + if isinstance(meta, Mapping): + # arm_pre_call writes a tuple; it arrives as a list once the metadata + # has been round-tripped through JSON. suppressed = meta.get(AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY) - if isinstance(suppressed, list) and marker in suppressed: + if isinstance(suppressed, (list, tuple)) and marker in suppressed: return True return False diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 534b2db3e61..5e6c9b34332 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -2009,7 +2009,7 @@ class ProxyBaseLLMRequestProcessing: # request: suppress every other compression guardrail and arm whichever one # the policy names for the model call, before those guardrails get a chance # to run below. - self.data = await _arm_auto_router_compression(data=self.data, llm_router=llm_router) + 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, diff --git a/litellm/proxy/guardrails/auto_router_compression.py b/litellm/proxy/guardrails/auto_router_compression.py index 2122d353def..d26479f7de0 100644 --- a/litellm/proxy/guardrails/auto_router_compression.py +++ b/litellm/proxy/guardrails/auto_router_compression.py @@ -13,8 +13,9 @@ each hop sees. """ import contextvars -from collections.abc import Mapping, MutableMapping, Sequence +from collections.abc import Iterable, Mapping, MutableMapping, Sequence from dataclasses import dataclass +from types import MappingProxyType from typing import TYPE_CHECKING, Final from litellm._logging import verbose_proxy_logger @@ -84,11 +85,11 @@ def policy_for_model( """ if llm_router is None: return None - deployments: Final = llm_router.get_model_list(model_name=model_alias, team_id=team_id) or [] + 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") or {}, Mapping) + if isinstance(litellm_params := deployment.get("litellm_params"), Mapping) and str(litellm_params.get("model", "")).startswith(AUTO_ROUTER_MODEL_PREFIX) ) requested: Final = frozenset(request_tags) @@ -128,7 +129,10 @@ def _active_compression_guardrails() -> tuple["CustomGuardrail", ...]: return tuple(cb for cb in active if isinstance(cb, compression_classes) and cb.guardrail_name) -async def arm_pre_call(data: MutableMapping[str, object], llm_router: "Router | None") -> MutableMapping[str, object]: +async def arm_pre_call( + data: MutableMapping[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, re-enables the model-side @@ -138,11 +142,11 @@ async def arm_pre_call(data: MutableMapping[str, object], llm_router: "Router | """ _routing_messages_snapshot.set(None) if llm_router is None: - return data + return model_alias: Final = data.get("model") if not isinstance(model_alias, str) or not model_alias: - return data + return # Read-only until a policy is confirmed: creating the metadata bucket for every # request, including the vast majority with no auto-router compression policy, @@ -156,7 +160,7 @@ async def arm_pre_call(data: MutableMapping[str, object], llm_router: "Router | request_tags=_get_tags_from_request_kwargs(data), ) if policy is None: - return data + return _, metadata = get_or_create_metadata_bucket(data) # Markers carry a per-process token so a caller cannot suppress a guardrail by @@ -167,35 +171,41 @@ async def arm_pre_call(data: MutableMapping[str, object], llm_router: "Router | if guardrail.guardrail_name != policy.model and (marker := guardrail.auto_router_suppression_marker()) ) if suppressed: - metadata[AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY] = list(suppressed) + metadata[AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY] = suppressed if policy.model is not None: - requested = metadata.get("guardrails") - if isinstance(requested, list): - if policy.model not in requested: - requested.append(policy.model) - else: - metadata["guardrails"] = [policy.model] + requested: Final = metadata.get("guardrails") + existing: Final = tuple(requested) if isinstance(requested, (list, tuple)) else () + if policy.model not in existing: + # A list, not a tuple: litellm_pre_call_utils tests this key with + # isinstance(..., list) and extends it, and would drop a tuple on the floor. + metadata["guardrails"] = [*existing, policy.model] # mutable-ok: this key's contract is a list from litellm.litellm_core_utils.prompt_templates.factory import resolve_structured_messages snapshot: Final = resolve_structured_messages(messages=data.get("messages"), request_kwargs=data) if snapshot is not None: - _routing_messages_snapshot.set(tuple(dict(message) for message in snapshot)) - - return data + _routing_messages_snapshot.set(tuple(MappingProxyType(dict(message)) for message in snapshot)) -def _snapshot_messages() -> list[dict[str, object]] | None: - snapshot: Final = _routing_messages_snapshot.get() - return None if snapshot is None else [dict(message) for message in snapshot] +def _snapshot_messages() -> tuple[Mapping[str, object], ...] | None: + return _routing_messages_snapshot.get() + + +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, - messages: list[dict[str, object]] | None, + # list[dict], not Sequence[Mapping]: the async_pre_routing_hook protocol in + # litellm/types/router.py types `messages` as list[dict[str, Any]]. + 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: +) -> 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`. Returns None when the caller should route on whatever messages it already has. @@ -207,12 +217,13 @@ async def messages_for_routing( if policy is None: return None - original: Final = _snapshot_messages() or messages + snapshot: Final = _snapshot_messages() + original: Final = snapshot if snapshot is not None else messages if policy.routing is None: # Explicitly no compression for routing. When the model side compressed, the # messages in hand are its output, so fall back to the untouched snapshot. - return _snapshot_messages() if policy.model is not None else None + return _as_routing_messages(snapshot) if policy.model is not None and snapshot is not None else None if not original: return None @@ -226,20 +237,20 @@ async def messages_for_routing( verbose_proxy_logger.warning( "AutoRouter compression: guardrail '%s' not found; routing on uncompressed messages", policy.routing ) - return original + return _as_routing_messages(original) - inputs: GenericGuardrailAPIInputs = { - "structured_messages": [dict(m) for m in original] # pyright: ignore[reportAssignmentType] # plain dicts, not AllMessageValues; see headroom.py's own use of this shape - } - # A throwaway request_data: apply_guardrail writes its stats onto this dict, not - # the real request's metadata, so routing-side compression never double-counts - # against extract_compression_saved_tokens's model-savings accounting. - throwaway_request_data: Final[dict[str, object]] = { - "messages": original, - "model": request_kwargs.get("model"), + inputs: Final[GenericGuardrailAPIInputs] = { + "structured_messages": _as_routing_messages(original) # pyright: ignore[reportAssignmentType] # plain dicts, not AllMessageValues; see headroom.py's own use of this shape } + model: Final = request_kwargs.get("model") + # A throwaway request_data: apply_guardrail writes its stats onto this dict, not the + # real request's metadata, so routing-side compression never double-counts against + # extract_compression_saved_tokens's model-savings accounting. + stats_sink: Final = {"messages": original, "model": model} # mutable-ok: apply_guardrail writes its stats here result: Final = await guardrail.apply_guardrail( - inputs=inputs, request_data=throwaway_request_data, input_type="request" + inputs=inputs, + request_data=stats_sink, + input_type="request", ) - compressed = result.get("structured_messages") - return compressed if isinstance(compressed, list) else original + compressed: Final = result.get("structured_messages") + return compressed if isinstance(compressed, list) else _as_routing_messages(original) diff --git a/litellm/router.py b/litellm/router.py index bcb2e2aa7ff..9637ad98c9a 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -13084,7 +13084,8 @@ class Router: and routing_messages is not None and pre_routing_hook_response.messages == routing_messages ): - pre_routing_hook_response = pre_routing_hook_response.model_copy(update={"messages": messages}) + restored: Final = {"messages": messages} # mutable-ok: pydantic's model_copy takes a dict + pre_routing_hook_response = pre_routing_hook_response.model_copy(update=restored) 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/tests/test_litellm/proxy/guardrails/test_auto_router_compression.py b/tests/test_litellm/proxy/guardrails/test_auto_router_compression.py index 79f069d8a2e..de667e8ed48 100644 --- a/tests/test_litellm/proxy/guardrails/test_auto_router_compression.py +++ b/tests/test_litellm/proxy/guardrails/test_auto_router_compression.py @@ -97,9 +97,7 @@ class TestPolicyForModel: 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"}}] - ) + 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): @@ -163,9 +161,7 @@ class _RecordingCompressionGuardrail(CustomGuardrail): ) -> 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 - ] + compressed = [{**m, "content": f"[COMPRESSED] {m.get('content')}"} for m in structured_messages] return {**inputs, "structured_messages": compressed} @@ -183,19 +179,16 @@ class TestArmPreCall: @pytest.mark.asyncio async def test_no_router_is_noop(self): data = {"model": "smart-router", "messages": [{"role": "user", "content": "hi"}]} - result = await arm_pre_call(data=data, llm_router=None) - assert result == data - assert "metadata" not in result + 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"}}] - ) + router = _FakeRouter([{"model_name": "smart-router", "litellm_params": {"model": "openai/gpt-4o-mini"}}]) data = {"model": "smart-router", "messages": [{"role": "user", "content": "hi"}]} - result = await arm_pre_call(data=data, llm_router=router) - assert "metadata" not in result - assert "litellm_metadata" not in result + 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): @@ -226,12 +219,12 @@ class TestArmPreCall: ] ) data = {"model": "smart-router", "messages": [{"role": "user", "content": "hi"}]} - result = await arm_pre_call(data=data, llm_router=router) - suppressed = result["metadata"][AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY] - assert suppressed == [always_on.auto_router_suppression_marker()] + await arm_pre_call(data=data, llm_router=router) + suppressed = data["metadata"][AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY] + assert tuple(suppressed) == (always_on.auto_router_suppression_marker(),) # The bare name alone must never suppress: that is what a caller could forge. assert "always-on-compression" not in suppressed - assert always_on.should_run_guardrail(data=result, event_type=GuardrailEventHooks.pre_call) is False + 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) @@ -262,8 +255,8 @@ class TestArmPreCall: ] ) data = {"model": "smart-router", "messages": [{"role": "user", "content": "hi"}]} - result = await arm_pre_call(data=data, llm_router=router) - assert result["metadata"]["guardrails"] == ["headroom-b"] + await arm_pre_call(data=data, llm_router=router) + assert data["metadata"]["guardrails"] == ["headroom-b"] @pytest.mark.asyncio async def test_snapshot_never_lands_in_persisted_metadata(self): @@ -275,10 +268,10 @@ class TestArmPreCall: original_messages = [{"role": "user", "content": "my ssn is 123-45-6789"}] data = {"model": "smart-router", "messages": original_messages} - result = await arm_pre_call(data=data, llm_router=router) + await arm_pre_call(data=data, llm_router=router) - assert "123-45-6789" not in json.dumps(result["metadata"]) - assert auto_router_compression._snapshot_messages() == original_messages + assert "123-45-6789" not in json.dumps(data["metadata"]) + assert [dict(m) for m in auto_router_compression._snapshot_messages()] == original_messages @pytest.mark.asyncio async def test_snapshot_is_a_copy_not_the_live_message_list(self): @@ -288,19 +281,19 @@ class TestArmPreCall: await arm_pre_call(data={"model": "smart-router", "messages": original_messages}, llm_router=router) original_messages[0]["content"] = "mutated after the snapshot" - assert auto_router_compression._snapshot_messages() == [{"role": "user", "content": "hi"}] + assert [dict(m) for m in auto_router_compression._snapshot_messages()] == [{"role": "user", "content": "hi"}] @pytest.mark.asyncio async def test_a_request_without_a_policy_clears_a_previous_snapshot(self): router_with = _FakeRouter([_marker({"auto_router_routing_compression": "headroom-a"})]) - await arm_pre_call(data={"model": "smart-router", "messages": [{"role": "user", "content": "first"}]}, - llm_router=router_with) - - router_without = _FakeRouter( - [{"model_name": "plain", "litellm_params": {"model": "openai/gpt-4o-mini"}}] + await arm_pre_call( + data={"model": "smart-router", "messages": [{"role": "user", "content": "first"}]}, llm_router=router_with + ) + + router_without = _FakeRouter([{"model_name": "plain", "litellm_params": {"model": "openai/gpt-4o-mini"}}]) + await arm_pre_call( + data={"model": "plain", "messages": [{"role": "user", "content": "second"}]}, llm_router=router_without ) - await arm_pre_call(data={"model": "plain", "messages": [{"role": "user", "content": "second"}]}, - llm_router=router_without) assert auto_router_compression._snapshot_messages() is None @@ -358,15 +351,11 @@ class TestMessagesForRouting: # rewrote `data["messages"]` to -- routing must ignore it and compress the # pristine snapshot instead. already_rewritten = [{"role": "user", "content": "rewritten by another guardrail"}] - result = await messages_for_routing( - policy=policy, messages=already_rewritten, request_kwargs={} - ) + result = await messages_for_routing(policy=policy, messages=already_rewritten, request_kwargs={}) assert result == [{"role": "user", "content": "[COMPRESSED] original"}] @pytest.mark.asyncio - async def test_guardrail_receives_a_throwaway_request_data_not_the_real_request_kwargs( - self, registered_guardrail - ): + async def test_guardrail_receives_a_throwaway_request_data_not_the_real_request_kwargs(self, registered_guardrail): """Regression: a real compression guardrail writes its stats onto whatever `request_data` dict it's given (`add_standard_logging_guardrail_information_to_ request_data`). If that were the caller's own `request_kwargs`, routing-side From 88ada40cdad9e711e7b40d5d174464667474585c Mon Sep 17 00:00:00 2001 From: moe-berri Date: Fri, 4 Sep 2026 18:37:38 -0700 Subject: [PATCH 06/13] fix(type-checking): satisfy the basedpyright budget gate for auto-router compression Two fixes for the zero-headroom basedpyright budget: - arm_pre_call's data parameter is dict[str, object], not MutableMapping: the latter is itself banned by LIT001 with no benefit, and it mismatched every dict-typed helper (get_or_create_metadata_bucket, resolve_structured_messages, _get_tags_from_request_kwargs), which is what the budget was actually flagging. - Router.async_pre_routing_hook computed pre_routing_hook_response in one shot instead of reassigning a Final-annotated local. The remaining two reportArgumentType hits are pre-existing: LiteLLM_Params(**merged) in _create_deployment_object already fails this check for all ~165 of its other fields, since the merged dict's value type is partly untyped/float; adding two new string fields to the model just grows that existing pile by two. Suppressed at the one call site with a reason, since fixing the root typing is out of scope here. --- .../proxy/guardrails/auto_router_compression.py | 12 ++++++++---- litellm/router.py | 16 +++++++--------- 2 files changed, 15 insertions(+), 13 deletions(-) diff --git a/litellm/proxy/guardrails/auto_router_compression.py b/litellm/proxy/guardrails/auto_router_compression.py index d26479f7de0..3b0804e40e2 100644 --- a/litellm/proxy/guardrails/auto_router_compression.py +++ b/litellm/proxy/guardrails/auto_router_compression.py @@ -13,7 +13,7 @@ each hop sees. """ import contextvars -from collections.abc import Iterable, Mapping, MutableMapping, Sequence +from collections.abc import Iterable, Mapping, Sequence from dataclasses import dataclass from types import MappingProxyType from typing import TYPE_CHECKING, Final @@ -89,7 +89,7 @@ def policy_for_model( markers: Final = tuple( litellm_params for deployment in deployments - if isinstance(litellm_params := deployment.get("litellm_params"), Mapping) + 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) @@ -130,7 +130,7 @@ def _active_compression_guardrails() -> tuple["CustomGuardrail", ...]: async def arm_pre_call( - data: MutableMapping[str, object], # mutable-ok: arms the live request dict in place + 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. @@ -183,7 +183,11 @@ async def arm_pre_call( from litellm.litellm_core_utils.prompt_templates.factory import resolve_structured_messages - snapshot: Final = resolve_structured_messages(messages=data.get("messages"), request_kwargs=data) + raw_messages: Final = data.get("messages") + snapshot: Final = resolve_structured_messages( + messages=raw_messages if isinstance(raw_messages, list) else None, + request_kwargs=data, + ) if snapshot is not None: _routing_messages_snapshot.set(tuple(MappingProxyType(dict(message)) for message in snapshot)) diff --git a/litellm/router.py b/litellm/router.py index 9637ad98c9a..989914b1610 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -8644,7 +8644,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}) @@ -13066,7 +13066,7 @@ class Router: else None ) - pre_routing_hook_response: Final = await selected_strategy.strategy.async_pre_routing_hook( + routed: Final = await selected_strategy.strategy.async_pre_routing_hook( model=registered_model_name, request_kwargs=request_kwargs, messages=routing_messages if routing_messages is not None else messages, @@ -13079,13 +13079,11 @@ class Router: # Compared by value, not identity: PreRoutingHookResponse is a pydantic model, # and pydantic reconstructs a validated list field rather than keeping the # exact object passed in, even when nothing about it changed. - if ( - pre_routing_hook_response is not None - and routing_messages is not None - and pre_routing_hook_response.messages == routing_messages - ): - restored: Final = {"messages": messages} # mutable-ok: pydantic's model_copy takes a dict - pre_routing_hook_response = pre_routing_hook_response.model_copy(update=restored) + 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), From 6385c7b3c53617fb480f5c8875a9b45538174ddf Mon Sep 17 00:00:00 2001 From: moe-berri Date: Fri, 4 Sep 2026 20:28:17 -0700 Subject: [PATCH 07/13] fix(auto-router compression): close three review findings on the per-hop policy Suppression state moves out of request metadata into a request-scoped ContextVar. refresh_proxy_server_request_body_snapshot copies metadata into proxy_server_request.body, which deployments persist to spend logs, so the marker naming each suppressed guardrail was readable by the caller whose request produced it. Recovering it was enough to replay {token}:{name} for any CustomGuardrail and switch off a PII or content-filter guardrail, since the check never verified the named guardrail was a compression one. Nothing is read from metadata now, so there is no marker to forge and the per-process token is no longer needed. Routing-side compression reads the live messages instead of a pre-guardrail copy. arm_pre_call runs before the pre-call hook, so its snapshot held the prompt as it was before any masking guardrail rewrote it, and messages_for_routing handed that to a compression guardrail which POSTs it to an external service. Masked content left the proxy anyway. The cost is one combination: when the model hop compressed and the hops differ, routing now classifies on the compressed text, since no uncompressed copy survives that a masking guardrail has already seen. policy_for_model no longer falls back to a marker scoped to tags the request does not carry, which applied an 'eu' policy to a 'us' request on config order alone. Each fix carries a regression test; all three fail when the fix is reverted. --- litellm/constants.py | 1 - litellm/integrations/custom_guardrail.py | 36 +- .../guardrails/auto_router_compression.py | 101 +++--- .../integrations/test_custom_guardrail.py | 335 +++++------------- .../test_auto_router_compression.py | 147 ++++---- tests/test_litellm/test_router.py | 25 +- 6 files changed, 236 insertions(+), 409 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 25fdaec20de..43fefaae048 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -217,7 +217,6 @@ PRE_CALL_EXECUTED_GUARDRAILS_KEY: Final = "_pre_call_executed_guardrails" # Metadata key listing compression guardrails an auto router's own compression # policy suppresses for this request. See litellm.proxy.guardrails.auto_router_compression. -AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY: Final = "_auto_router_suppressed_compression_guardrails" # Generic fallback for unknown models DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET: Final = int( diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 1ec08641706..f511d128dfc 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -45,7 +45,6 @@ dc: Final = DualCache() from litellm.constants import ( - AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY, GUARDRAIL_SCANNED_MESSAGES_CACHE_TTL_SECONDS, PRE_CALL_EXECUTED_GUARDRAILS_KEY, ) @@ -941,33 +940,22 @@ class CustomGuardrail(CustomLogger): """ return False - def auto_router_suppression_marker(self) -> str | None: - """The value `arm_pre_call` must write to suppress this guardrail. + def _suppressed_by_auto_router_compression(self) -> bool: + """True when an auto router's own compression policy suppresses this guardrail. - Carries the per-process token for the same reason `_pre_call_marker` does: a - caller controls request metadata, so a bare guardrail name there would let any - request switch off a PII, content-filter, or compression guardrail for itself. - The token is never sent to the caller, so the marker cannot be forged. + 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 None - return f"{_PRE_CALL_EXECUTED_TOKEN}:{name}" - - def _suppressed_by_auto_router_compression(self, data: Mapping[str, object]) -> bool: - """True when an auto router's own compression policy suppresses this guardrail.""" - marker: Final = self.auto_router_suppression_marker() - if marker is None: return False - for meta_key in ("metadata", "litellm_metadata"): - meta = data.get(meta_key) - if isinstance(meta, Mapping): - # arm_pre_call writes a tuple; it arrives as a list once the metadata - # has been round-tripped through JSON. - suppressed = meta.get(AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY) - if isinstance(suppressed, (list, tuple)) and marker in suppressed: - return True - return False + from litellm.proxy.guardrails.auto_router_compression import ( + suppressed_compression_guardrails, + ) + + return name in suppressed_compression_guardrails() def should_run_guardrail( self, @@ -977,7 +965,7 @@ class CustomGuardrail(CustomLogger): """ Returns True if the guardrail should be run on the event_type """ - if self._suppressed_by_auto_router_compression(data): + if self._suppressed_by_auto_router_compression(): return False requested_guardrails: Final = self.get_guardrail_from_metadata(data) diff --git a/litellm/proxy/guardrails/auto_router_compression.py b/litellm/proxy/guardrails/auto_router_compression.py index 3b0804e40e2..2a4a0c82022 100644 --- a/litellm/proxy/guardrails/auto_router_compression.py +++ b/litellm/proxy/guardrails/auto_router_compression.py @@ -15,11 +15,9 @@ each hop sees. import contextvars from collections.abc import Iterable, Mapping, Sequence from dataclasses import dataclass -from types import MappingProxyType from typing import TYPE_CHECKING, Final from litellm._logging import verbose_proxy_logger -from litellm.constants import AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY from litellm.litellm_core_utils.core_helpers import get_or_create_metadata_bucket from litellm.router_utils.auto_router_model_naming import AUTO_ROUTER_MODEL_PREFIX from litellm.types.utils import GenericGuardrailAPIInputs @@ -31,16 +29,22 @@ if TYPE_CHECKING: COMPRESSION_GUARDRAIL_PROVIDERS: Final = frozenset({"headroom", "compresr"}) _NO_COMPRESSION: Final = "none" -# The pre-compression messages, so a routing decision that does not share the model -# call's compression still classifies on the original text. Deliberately a ContextVar -# rather than a metadata key: `refresh_proxy_server_request_body_snapshot` copies -# metadata into `proxy_server_request.body`, which deployments persist to spend logs, -# and this holds the prompt as it was before any masking guardrail rewrote it. -_routing_messages_snapshot: Final[contextvars.ContextVar[tuple[Mapping[str, object], ...] | None]] = ( - contextvars.ContextVar("litellm_auto_router_routing_messages_snapshot", default=None) +# Compression guardrails this request's auto router has switched off. Deliberately a +# ContextVar rather than a metadata key: `refresh_proxy_server_request_body_snapshot` +# copies metadata into `proxy_server_request.body`, which deployments persist to spend +# logs. A suppression list that reaches a log the caller can read is a list the caller +# can replay, which would let any request switch off a PII or content-filter guardrail. +# Nothing here is caller-supplied, so there is no marker to forge in the first place. +_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() + + @dataclass(frozen=True, slots=True) class AutoRouterCompressionPolicy: """An auto router's compression choice for each hop. ``None`` means no compression.""" @@ -96,7 +100,11 @@ def policy_for_model( tag_matched: Final = tuple( params for params in markers if (tags := params.get("tags")) and requested.issuperset(frozenset(tags)) ) - for params in (*tag_matched, *markers): + # Only untagged markers may serve as the fallback. A marker scoped to tags this + # request does not carry describes a different slice of traffic, so falling back + # to it would apply, say, an "eu" policy to a "us" request purely on config order. + untagged: Final = tuple(params for params in markers if not params.get("tags")) + for params in (*tag_matched, *untagged): policy = policy_from_litellm_params(params) if policy is not None: return policy @@ -135,12 +143,10 @@ async def arm_pre_call( ) -> None: """Apply an auto router's compression policy, if any, before guardrails run. - Suppresses every other compression guardrail, re-enables the model-side - guardrail the policy names (if any) even when it isn't ``default_on``, and - snapshots the pre-compression messages so the routing decision can read them - independently of whatever the model-side guardrail does to `data`. + Suppresses every other compression guardrail and re-enables the model-side + guardrail the policy names (if any) even when it isn't ``default_on``. """ - _routing_messages_snapshot.set(None) + _suppressed_compression_guardrails.set(frozenset()) if llm_router is None: return @@ -162,18 +168,16 @@ async def arm_pre_call( if policy is None: return - _, metadata = get_or_create_metadata_bucket(data) - # Markers carry a per-process token so a caller cannot suppress a guardrail by - # naming it in its own request metadata. - suppressed: Final = tuple( - marker - for guardrail in _active_compression_guardrails() - if guardrail.guardrail_name != policy.model and (marker := guardrail.auto_router_suppression_marker()) + _suppressed_compression_guardrails.set( + frozenset( + name + for guardrail in _active_compression_guardrails() + if (name := guardrail.guardrail_name) and name != policy.model + ) ) - if suppressed: - metadata[AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY] = suppressed if policy.model is not None: + _, 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: @@ -181,20 +185,6 @@ async def arm_pre_call( # isinstance(..., list) and extends it, and would drop a tuple on the floor. metadata["guardrails"] = [*existing, policy.model] # mutable-ok: this key's contract is a list - from litellm.litellm_core_utils.prompt_templates.factory import resolve_structured_messages - - raw_messages: Final = data.get("messages") - snapshot: Final = resolve_structured_messages( - messages=raw_messages if isinstance(raw_messages, list) else None, - request_kwargs=data, - ) - if snapshot is not None: - _routing_messages_snapshot.set(tuple(MappingProxyType(dict(message)) for message in snapshot)) - - -def _snapshot_messages() -> tuple[Mapping[str, object], ...] | None: - return _routing_messages_snapshot.get() - def _as_routing_messages( messages: Iterable[Mapping[str, object]], @@ -213,23 +203,22 @@ async def messages_for_routing( """Messages to use for a routing decision, per `policy.routing`. Returns None when the caller should route on whatever messages it already has. - The model call is untouched either way: model-side compression, if any, already - ran as an ordinary pre-call guardrail before the router was reached, so when the - two hops differ the routing decision reads the pre-compression snapshot rather - than what that guardrail left behind. + + Always reads the live messages, never a pre-guardrail copy of them. The routing + hop compresses through a real guardrail, which POSTs the text to an external + compression service, so it must see what every other guardrail has already done + to the request. Routing on a snapshot taken before the pre-call hook would send + a masking guardrail's own input straight back out of the proxy. + + The consequence, when the model hop compressed and the two hops differ: the + messages in hand are that guardrail's output, and there is no un-compressed copy + left to route on. The routing decision reads the compressed text in that one + combination rather than leaking the original. """ - if policy is None: + if policy is None or policy.routing is None: return None - snapshot: Final = _snapshot_messages() - original: Final = snapshot if snapshot is not None else messages - - if policy.routing is None: - # Explicitly no compression for routing. When the model side compressed, the - # messages in hand are its output, so fall back to the untouched snapshot. - return _as_routing_messages(snapshot) if policy.model is not None and snapshot is not None else None - - if not original: + if not messages: return None from litellm.proxy.common_utils.registry_read_through import ( @@ -241,20 +230,20 @@ async def messages_for_routing( verbose_proxy_logger.warning( "AutoRouter compression: guardrail '%s' not found; routing on uncompressed messages", policy.routing ) - return _as_routing_messages(original) + return _as_routing_messages(messages) inputs: Final[GenericGuardrailAPIInputs] = { - "structured_messages": _as_routing_messages(original) # pyright: ignore[reportAssignmentType] # plain dicts, not AllMessageValues; see headroom.py's own use of this shape + "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") # A throwaway request_data: apply_guardrail writes its stats onto this dict, not the # real request's metadata, so routing-side compression never double-counts against # extract_compression_saved_tokens's model-savings accounting. - stats_sink: Final = {"messages": original, "model": model} # mutable-ok: apply_guardrail writes its stats here + 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(original) + return compressed if isinstance(compressed, list) else _as_routing_messages(messages) diff --git a/tests/test_litellm/integrations/test_custom_guardrail.py b/tests/test_litellm/integrations/test_custom_guardrail.py index f590903cb74..c4a702d453c 100644 --- a/tests/test_litellm/integrations/test_custom_guardrail.py +++ b/tests/test_litellm/integrations/test_custom_guardrail.py @@ -13,7 +13,6 @@ from litellm.types.utils import GenericGuardrailAPIInputs, GuardrailTracingDetai 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""" @@ -26,18 +25,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 @@ -64,9 +59,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() @@ -99,9 +92,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 @@ -114,9 +105,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 @@ -130,9 +119,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 @@ -144,9 +131,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 @@ -175,9 +160,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 @@ -189,15 +172,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 @@ -214,9 +194,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 @@ -236,9 +214,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 @@ -255,9 +231,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 @@ -277,9 +251,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 @@ -298,9 +270,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 @@ -312,9 +282,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 = { @@ -345,12 +313,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. @@ -361,12 +325,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, @@ -432,12 +392,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 = { @@ -446,10 +401,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 @@ -458,12 +410,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 = { @@ -471,12 +418,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 = { @@ -485,10 +427,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): @@ -511,17 +450,12 @@ 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.constants import AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY + from litellm.proxy.guardrails import auto_router_compression from litellm.types.guardrails import GuardrailEventHooks always_on = CustomGuardrail( @@ -529,22 +463,17 @@ class TestCustomGuardrailShouldRunGuardrail: default_on=True, event_hook=GuardrailEventHooks.pre_call, ) - data = { - "model": "smart-router", - "metadata": { - AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY: [ - always_on.auto_router_suppression_marker() - ], - }, - } + 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) - assert ( - always_on.should_run_guardrail(data=data, event_type=GuardrailEventHooks.pre_call) - is False - ) - - def test_should_run_guardrail_suppression_list_does_not_affect_other_names(self): - from litellm.constants import AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY + 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( @@ -552,25 +481,20 @@ class TestCustomGuardrailShouldRunGuardrail: default_on=True, event_hook=GuardrailEventHooks.pre_call, ) - other = CustomGuardrail(guardrail_name="some-other-guardrail", default_on=True) - data = { - "model": "smart-router", - "metadata": { - AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY: [ - other.auto_router_suppression_marker() - ], - }, - } + 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) - assert ( - always_on.should_run_guardrail(data=data, event_type=GuardrailEventHooks.pre_call) - is True - ) - - def test_should_run_guardrail_ignores_a_forged_suppression_marker(self): - """A caller controls request metadata, so a bare guardrail name there must not - switch off an always-on guardrail: only the per-process marker counts.""" - from litellm.constants import AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY + 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( @@ -581,17 +505,14 @@ class TestCustomGuardrailShouldRunGuardrail: forged = { "model": "smart-router", "metadata": { - AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY: [ + "_auto_router_suppressed_compression_guardrails": [ "headroom-default", - "forged-token:headroom-default", + "any-token:headroom-default", ], }, } - assert ( - always_on.should_run_guardrail(data=forged, event_type=GuardrailEventHooks.pre_call) - is True - ) + assert always_on.should_run_guardrail(data=forged, event_type=GuardrailEventHooks.pre_call) is True class TestApplyGuardrailCheck: @@ -630,35 +551,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: @@ -685,11 +604,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) @@ -701,11 +616,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) @@ -717,18 +628,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" @@ -745,12 +650,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" @@ -768,9 +671,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 @@ -816,9 +717,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 @@ -828,9 +727,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) @@ -927,9 +824,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) @@ -942,9 +837,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"}, ], @@ -998,9 +891,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 @@ -1009,10 +900,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 @@ -1071,9 +959,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", ) @@ -1096,9 +982,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 @@ -1110,18 +994,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:]) @@ -1615,9 +1495,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, ( @@ -1659,9 +1537,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 @@ -1790,9 +1666,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"] @@ -1834,13 +1708,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": {}} @@ -1851,17 +1719,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 index de667e8ed48..0b47e56cb02 100644 --- a/tests/test_litellm/proxy/guardrails/test_auto_router_compression.py +++ b/tests/test_litellm/proxy/guardrails/test_auto_router_compression.py @@ -4,15 +4,16 @@ Unit tests for litellm.proxy.guardrails.auto_router_compression. Covers: - policy_from_litellm_params: absent keys mean no policy; the "none" sentinel normalizes to explicit no-compression within an active policy; is_same -- policy_for_model: finds the auto-router marker deployment for an alias, and - picks the tag-scoped marker the request's tags actually match +- policy_for_model: finds the auto-router marker deployment for an alias, picks the + tag-scoped marker the request's tags actually match, and never falls back to a + marker scoped to tags the request does not carry - arm_pre_call: no-op without a policy; suppresses active compression guardrails - with a forgery-proof marker; arms the model-side guardrail even when it isn't - default_on; keeps the pre-compression snapshot out of persisted metadata -- messages_for_routing: no-op without a policy; routes on the pre-compression - snapshot when the two hops differ; compresses via the named guardrail's - apply_guardrail; never writes stats onto the caller's own request_kwargs - (regression for double-counted compression savings) + through request-scoped state rather than metadata, which reaches spend logs a + caller can read; arms the model-side guardrail even when it isn't default_on +- messages_for_routing: no-op without a policy; compresses the live messages every + earlier guardrail has already rewritten, never a pre-guardrail copy of them; + never writes stats onto the caller's own request_kwargs (regression for + double-counted compression savings) """ import json @@ -20,7 +21,6 @@ from typing import Any import pytest -from litellm.constants import AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.proxy.guardrails import auto_router_compression from litellm.proxy.guardrails.auto_router_compression import ( @@ -136,6 +136,26 @@ class TestPolicyForModel: ) assert policy == AutoRouterCompressionPolicy(routing="headroom-a", model=None) + def test_a_marker_scoped_to_other_tags_is_never_the_fallback(self): + """Regression: an "eu" marker describes a different slice of traffic, so a "us" + request must not fall back to its policy just because it is configured first.""" + 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): + """With only tag-scoped markers and none matching, there is no policy to apply: + inheriting an unrelated slice's compression is worse than inheriting nothing.""" + 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.""" @@ -220,25 +240,30 @@ class TestArmPreCall: ) data = {"model": "smart-router", "messages": [{"role": "user", "content": "hi"}]} await arm_pre_call(data=data, llm_router=router) - suppressed = data["metadata"][AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY] - assert tuple(suppressed) == (always_on.auto_router_suppression_marker(),) - # The bare name alone must never suppress: that is what a caller could forge. - assert "always-on-compression" not in suppressed + 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_a_caller_cannot_suppress_a_guardrail_by_naming_it_in_metadata(self): - """Regression: request metadata is caller-controlled, so a bare guardrail name - there must not switch off a PII, content-filter, or compression guardrail.""" + async def test_suppression_state_never_enters_request_metadata(self): + """Regression (security): a suppression list written to metadata is copied into + proxy_server_request.body and persisted to spend logs, so a caller could read it + back and replay it to switch off a PII or content-filter guardrail.""" guardrail = _RecordingCompressionGuardrail(guardrail_name="always-on-compression") - forged = { - "model": "smart-router", - "metadata": {AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY: ["always-on-compression"]}, - } + import litellm - assert guardrail._suppressed_by_auto_router_compression(forged) is False + 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): @@ -259,52 +284,20 @@ class TestArmPreCall: assert data["metadata"]["guardrails"] == ["headroom-b"] @pytest.mark.asyncio - async def test_snapshot_never_lands_in_persisted_metadata(self): - """Regression: refresh_proxy_server_request_body_snapshot copies metadata into - proxy_server_request.body, which deployments persist to spend logs. The - pre-compression snapshot holds the prompt before any masking guardrail ran, so - it must live outside anything that gets serialized.""" + async def test_arm_pre_call_keeps_no_copy_of_the_prompt(self): + """Regression (security): arm_pre_call runs before the pre-call guardrails, so + any copy of the messages it retained would be the pre-masking text. Routing-side + compression POSTs its input to an external service, so that copy must not exist.""" router = _FakeRouter([_marker({"auto_router_routing_compression": "headroom-a"})]) - original_messages = [{"role": "user", "content": "my ssn is 123-45-6789"}] - data = {"model": "smart-router", "messages": original_messages} + 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["metadata"]) - assert [dict(m) for m in auto_router_compression._snapshot_messages()] == original_messages - - @pytest.mark.asyncio - async def test_snapshot_is_a_copy_not_the_live_message_list(self): - router = _FakeRouter([_marker({"auto_router_routing_compression": "headroom-a"})]) - original_messages = [{"role": "user", "content": "hi"}] - - await arm_pre_call(data={"model": "smart-router", "messages": original_messages}, llm_router=router) - original_messages[0]["content"] = "mutated after the snapshot" - - assert [dict(m) for m in auto_router_compression._snapshot_messages()] == [{"role": "user", "content": "hi"}] - - @pytest.mark.asyncio - async def test_a_request_without_a_policy_clears_a_previous_snapshot(self): - router_with = _FakeRouter([_marker({"auto_router_routing_compression": "headroom-a"})]) - await arm_pre_call( - data={"model": "smart-router", "messages": [{"role": "user", "content": "first"}]}, llm_router=router_with - ) - - router_without = _FakeRouter([{"model_name": "plain", "litellm_params": {"model": "openai/gpt-4o-mini"}}]) - await arm_pre_call( - data={"model": "plain", "messages": [{"role": "user", "content": "second"}]}, llm_router=router_without - ) - - assert auto_router_compression._snapshot_messages() is None + assert "123-45-6789" not in json.dumps(data.get("metadata", {})) + assert not hasattr(auto_router_compression, "_routing_messages_snapshot") class TestMessagesForRouting: - @pytest.fixture(autouse=True) - def _clear_snapshot(self): - auto_router_compression._routing_messages_snapshot.set(None) - yield - auto_router_compression._routing_messages_snapshot.set(None) - @pytest.mark.asyncio async def test_no_policy_returns_none(self): assert await messages_for_routing(policy=None, messages=[], request_kwargs={}) is None @@ -316,18 +309,15 @@ class TestMessagesForRouting: assert await messages_for_routing(policy=policy, messages=[], request_kwargs={}) is None @pytest.mark.asyncio - async def test_routing_none_with_model_compression_routes_on_the_snapshot(self): - """Regression: with routing explicitly off and the model side compressed, the - messages in hand are the model-side guardrail's output. Routing asked for no - compression, so it must read the pre-compression snapshot instead.""" - original = [{"role": "user", "content": "the full original conversation"}] - auto_router_compression._routing_messages_snapshot.set(tuple(dict(m) for m in original)) + async def test_routing_none_never_reaches_for_a_pre_guardrail_copy(self): + """Routing asked for no compression while the model hop compressed, so the + messages in hand are that guardrail's output and no uncompressed copy survives. + Routing reads them as-is: the alternative is keeping a pre-guardrail copy, which + is the text a masking guardrail exists to remove.""" policy = AutoRouterCompressionPolicy(routing=None, model="headroom-a") model_compressed = [{"role": "user", "content": "[COMPRESSED] the full original conversation"}] - result = await messages_for_routing(policy=policy, messages=model_compressed, request_kwargs={}) - - assert result == original + 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): @@ -344,15 +334,18 @@ class TestMessagesForRouting: assert result == [{"role": "user", "content": "[COMPRESSED] hello world"}] @pytest.mark.asyncio - async def test_uses_the_snapshot_when_present(self, registered_guardrail): + async def test_routing_compresses_what_the_other_guardrails_left_behind(self, registered_guardrail): + """Regression (security): routing-side compression POSTs its input to an external + service, so it must read the live messages every earlier guardrail has already + rewritten. Routing on a pre-guardrail copy would send a masking guardrail's own + input straight back out of the proxy.""" policy = AutoRouterCompressionPolicy(routing="fake-compress", model="headroom-b") - auto_router_compression._routing_messages_snapshot.set(({"role": "user", "content": "original"},)) - # `messages` here stands in for whatever the model-side guardrail already - # rewrote `data["messages"]` to -- routing must ignore it and compress the - # pristine snapshot instead. - already_rewritten = [{"role": "user", "content": "rewritten by another guardrail"}] - result = await messages_for_routing(policy=policy, messages=already_rewritten, request_kwargs={}) - assert result == [{"role": "user", "content": "[COMPRESSED] original"}] + 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_guardrail_receives_a_throwaway_request_data_not_the_real_request_kwargs(self, registered_guardrail): diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 4c5813911ac..b4de20e9f0a 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -10106,32 +10106,29 @@ class TestAutoRouterCompressionDecoupling: assert registered_guardrail.call_count == 0 @pytest.mark.asyncio - async def test_routing_none_routes_on_the_original_not_the_model_compressed_messages( + async def test_routing_none_classifies_on_the_live_messages_not_a_pre_guardrail_copy( self, registered_guardrail ): - """Regression: with routing explicitly off and the model side compressed, the - messages the router holds are the model-side guardrail's output. Routing asked - for no compression, so it has to classify on the pre-compression snapshot.""" - from litellm.proxy.guardrails import auto_router_compression + """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", } ) - original_messages = self._messages() - auto_router_compression._routing_messages_snapshot.set(tuple(dict(m) for m in original_messages)) model_compressed = [{"role": "user", "content": "[COMPRESSED] What is the capital of France?"}] - try: - await router.async_pre_routing_hook( - model="smart-router", request_kwargs={"metadata": {}}, messages=model_compressed - ) - finally: - auto_router_compression._routing_messages_snapshot.set(None) + await router.async_pre_routing_hook( + model="smart-router", request_kwargs={"metadata": {}}, messages=model_compressed + ) - assert strategy.received_messages == original_messages + assert strategy.received_messages == model_compressed assert registered_guardrail.call_count == 0 @pytest.mark.asyncio From 00b49ccc8bb0de73891376e5ee1e8bd58295ba2d Mon Sep 17 00:00:00 2001 From: moe-berri Date: Sat, 5 Sep 2026 09:31:18 -0700 Subject: [PATCH 08/13] fix(auto-router compression): honor the policy on the SDK path and on re-save The router reused the model hop's compression for routing whenever both hops named the same guardrail, on the premise that arm_pre_call had already run it. Only the proxy calls arm_pre_call, so through the SDK nothing armed the guardrail and nothing had compressed anything: the shortcut skipped routing compression too and served the request with no compression on either hop. The reuse is now conditional on the model hop actually having been armed. The Admin UI hydrated an absent auto_router_model_compression as same-as-routing, while the backend reads it as no model-hop compression. Opening a router configured with only auto_router_routing_compression and saving any unrelated edit wrote the routing guardrail onto the model hop, silently starting to compress the model call. Both carry a regression test that fails when the fix is reverted. --- .../guardrails/auto_router_compression.py | 15 + litellm/router.py | 8 +- tests/test_litellm/test_router.py | 967 ++++++------------ .../buildAutoRouterCompression.test.ts | 15 +- .../add_model/buildAutoRouterCompression.ts | 8 +- 5 files changed, 347 insertions(+), 666 deletions(-) diff --git a/litellm/proxy/guardrails/auto_router_compression.py b/litellm/proxy/guardrails/auto_router_compression.py index 2ab779658b5..e7f58662249 100644 --- a/litellm/proxy/guardrails/auto_router_compression.py +++ b/litellm/proxy/guardrails/auto_router_compression.py @@ -45,6 +45,19 @@ def suppressed_compression_guardrails() -> frozenset[str]: return _suppressed_compression_guardrails.get() +# Whether `arm_pre_call` actually armed a model-side compression guardrail for this +# request. 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.""" @@ -147,6 +160,7 @@ async def arm_pre_call( 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 @@ -179,6 +193,7 @@ async def arm_pre_call( ) if policy.model is not None: + _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 () diff --git a/litellm/router.py b/litellm/router.py index 989914b1610..b61e4e29e09 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -13039,6 +13039,7 @@ class Router: from litellm.proxy.guardrails.auto_router_compression import ( messages_for_routing, + model_hop_compression_armed, policy_for_model, team_id_from_request, ) @@ -13057,8 +13058,13 @@ class Router: # (arm_pre_call armed it whether or not it is `default_on`); reuse that result # for routing too instead of paying for a second compression call against the # same content. + # + # Only the proxy calls arm_pre_call, so that reuse is conditional on it having + # actually run: on the SDK path nothing arms the model hop and nothing has + # compressed anything, and taking the shortcut there would skip both hops and + # silently serve the request with no compression at all. needs_independent_routing_compression: Final = compression_policy is not None and not ( - compression_policy.is_same and compression_policy.model is not None + 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) diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index a59c9c98163..9e6a88d3433 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -15,7 +15,6 @@ import pytest import respx - import litellm from litellm import Router from litellm.exceptions import MidStreamFallbackError @@ -137,31 +136,18 @@ def test_router_model_group_encrypted_content_affinity_callback_registration(): num_retries=0, ) callbacks = router.optional_callbacks or [] - encrypted_content_callbacks = [ - cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck) - ] - deployment_callback = next( - cb for cb in callbacks if isinstance(cb, DeploymentAffinityCheck) - ) + encrypted_content_callbacks = [cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck)] + deployment_callback = next(cb for cb in callbacks if isinstance(cb, DeploymentAffinityCheck)) assert len(encrypted_content_callbacks) == 1 assert encrypted_content_callbacks[0].enable_global_affinity is False - assert ( - encrypted_content_callbacks[0].model_group_affinity_config - == model_group_affinity_config - ) - assert callbacks.index(encrypted_content_callbacks[0]) < callbacks.index( - deployment_callback - ) - assert litellm.callbacks.index(encrypted_content_callbacks[0]) < ( - litellm.callbacks.index(deployment_callback) - ) + assert encrypted_content_callbacks[0].model_group_affinity_config == model_group_affinity_config + assert callbacks.index(encrypted_content_callbacks[0]) < callbacks.index(deployment_callback) + assert litellm.callbacks.index(encrypted_content_callbacks[0]) < (litellm.callbacks.index(deployment_callback)) router._add_encrypted_content_affinity_check(enable_global_affinity=True) callbacks = router.optional_callbacks or [] - encrypted_content_callbacks = [ - cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck) - ] + encrypted_content_callbacks = [cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck)] assert len(encrypted_content_callbacks) == 1 assert encrypted_content_callbacks[0].enable_global_affinity is True assert encrypted_content_callbacks[0].router is router @@ -192,13 +178,9 @@ async def test_encrypted_content_affinity_model_group_config_is_additive(): }, target_deployment, ] - encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id( - "deployment-b", "rs_test" - ) + encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id("deployment-b", "rs_test") - assert EncryptedContentAffinityCheck.has_model_group_affinity_enabled( - {model_group: ["encrypted_content_affinity"]} - ) + assert EncryptedContentAffinityCheck.has_model_group_affinity_enabled({model_group: ["encrypted_content_affinity"]}) assert not EncryptedContentAffinityCheck.has_model_group_affinity_enabled(None) per_group_check = EncryptedContentAffinityCheck( @@ -239,10 +221,7 @@ async def test_encrypted_content_affinity_model_group_config_is_additive(): ) assert unfiltered == healthy_deployments - assert ( - "encrypted_content_affinity_enabled" - not in disabled_request_kwargs["litellm_metadata"] - ) + assert "encrypted_content_affinity_enabled" not in disabled_request_kwargs["litellm_metadata"] global_check = EncryptedContentAffinityCheck( enable_global_affinity=True, @@ -262,9 +241,7 @@ async def test_encrypted_content_affinity_model_group_config_is_additive(): ) assert globally_filtered == [target_deployment] - assert global_request_kwargs["litellm_metadata"][ - "encrypted_content_affinity_enabled" - ] + assert global_request_kwargs["litellm_metadata"]["encrypted_content_affinity_enabled"] @pytest.mark.asyncio @@ -311,18 +288,10 @@ async def test_encrypted_content_affinity_takes_priority_over_user_key_affinity( num_retries=0, ) callbacks = router.optional_callbacks or [] - deployment_callback = next( - cb for cb in callbacks if isinstance(cb, DeploymentAffinityCheck) - ) - encrypted_content_callback = next( - cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck) - ) - assert callbacks.index(encrypted_content_callback) < callbacks.index( - deployment_callback - ) - assert litellm.callbacks.index(encrypted_content_callback) < ( - litellm.callbacks.index(deployment_callback) - ) + deployment_callback = next(cb for cb in callbacks if isinstance(cb, DeploymentAffinityCheck)) + encrypted_content_callback = next(cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck)) + assert callbacks.index(encrypted_content_callback) < callbacks.index(deployment_callback) + assert litellm.callbacks.index(encrypted_content_callback) < (litellm.callbacks.index(deployment_callback)) cache_key = DeploymentAffinityCheck.get_affinity_cache_key( model_group=model_group, @@ -333,9 +302,7 @@ async def test_encrypted_content_affinity_takes_priority_over_user_key_affinity( value={"model_id": "deployment-a"}, ttl=60, ) - encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id( - "deployment-b", "rs_test" - ) + encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id("deployment-b", "rs_test") request_kwargs = { "input": [{"type": "reasoning", "id": encoded_id}], "litellm_metadata": {"user_api_key_hash": user_api_key_hash}, @@ -939,9 +906,7 @@ async def test_arouter_aretrieve_batch(): ], ) - with patch.object( - litellm, "aretrieve_batch", return_value=AsyncMock() - ) as mock_aretrieve_batch: + with patch.object(litellm, "aretrieve_batch", return_value=AsyncMock()) as mock_aretrieve_batch: try: response = await router.aretrieve_batch( model="gpt-3.5-turbo", @@ -962,9 +927,7 @@ async def test_arouter_aretrieve_file_content(): Test that router.acreate_file with JSONL file returns the correct response """ - with patch.object( - litellm, "afile_content", return_value=AsyncMock() - ) as mock_afile_content: + with patch.object(litellm, "afile_content", return_value=AsyncMock()) as mock_afile_content: router = litellm.Router( model_list=[ { @@ -1025,7 +988,7 @@ async def test_arouter_filter_team_based_models(): assert result is not None # FAILS - with pytest.raises(Exception, match='No deployments available for selected model, Try again in') as e: + with pytest.raises(Exception, match="No deployments available for selected model, Try again in") as e: result = await router.acompletion( model="gpt-3.5-turbo", messages=[{"role": "user", "content": "Hello, world!"}], @@ -1109,9 +1072,7 @@ def test_arouter_should_include_deployment(): model=deployment_with_team_and_public_name, team_id="test-team", ) - assert ( - result is True - ), "Should return True when team_id and team_public_model_name match" + assert result is True, "Should return True when team_id and team_public_model_name match" # Test Case 2: Team-specific deployment - team_id matches but model_name doesn't match team_public_model_name result = router.should_include_deployment( @@ -1119,9 +1080,9 @@ def test_arouter_should_include_deployment(): model=deployment_with_team_and_public_name, team_id="test-team", ) - assert ( - result is False - ), "Should return False when team_id matches but model_name doesn't match team_public_model_name" + assert result is False, ( + "Should return False when team_id matches but model_name doesn't match team_public_model_name" + ) # Test Case 3: Team-specific deployment - team_id doesn't match result = router.should_include_deployment( @@ -1137,30 +1098,18 @@ def test_arouter_should_include_deployment(): model=deployment_with_team_no_public_name, team_id="test-team", ) - assert ( - result is True - ), "Should return True when team deployment has no team_public_model_name to match" + assert result is True, "Should return True when team deployment has no team_public_model_name to match" # Test Case 5: Non-team deployment - model_name matches and no team_id - result = router.should_include_deployment( - model_name="gpt-4", model=deployment_without_team, team_id=None - ) - assert ( - result is True - ), "Should return True when model_name matches and deployment has no team_id" + result = router.should_include_deployment(model_name="gpt-4", model=deployment_without_team, team_id=None) + assert result is True, "Should return True when model_name matches and deployment has no team_id" # Test Case 6: Non-team deployment - model_name matches but team_id provided (should still work) - result = router.should_include_deployment( - model_name="gpt-4", model=deployment_without_team, team_id="any-team" - ) - assert ( - result is True - ), "Should return True when model_name matches non-team deployment, regardless of team_id param" + result = router.should_include_deployment(model_name="gpt-4", model=deployment_without_team, team_id="any-team") + assert result is True, "Should return True when model_name matches non-team deployment, regardless of team_id param" # Test Case 7: Non-team deployment - model_name doesn't match - result = router.should_include_deployment( - model_name="different-model", model=deployment_without_team, team_id=None - ) + result = router.should_include_deployment(model_name="different-model", model=deployment_without_team, team_id=None) assert result is False, "Should return False when model_name doesn't match" # Test Case 8: Team deployment accessed without matching team_id @@ -1169,9 +1118,7 @@ def test_arouter_should_include_deployment(): model=deployment_with_team_and_public_name, team_id=None, ) - assert ( - result is True - ), "Should return True when matching model with exact model_name" + assert result is True, "Should return True when matching model with exact model_name" def test_arouter_responses_api_bridge(): @@ -1221,9 +1168,7 @@ def test_arouter_responses_api_bridge(): "status": "completed", "output": [], } - mock_response.text = ( - '{"id": "resp_test", "object": "response", "status": "completed", "output": []}' - ) + mock_response.text = '{"id": "resp_test", "object": "response", "status": "completed", "output": []}' with patch.object(client, "post", return_value=mock_response) as mock_post: try: @@ -1297,7 +1242,7 @@ def test_add_invalid_provider_to_router(): ], ) - with pytest.raises(Exception, match='Unsupported provider - vertex_ai_eu') as e: + with pytest.raises(Exception, match="Unsupported provider - vertex_ai_eu") as e: router.add_deployment( Deployment( model_name="vertex_ai/*", @@ -1349,15 +1294,9 @@ async def test_router_ageneric_api_call_with_fallbacks_helper(): }, } - with patch.object( - router, "_update_kwargs_with_deployment" - ) as mock_update_kwargs: - with patch.object( - router, "async_routing_strategy_pre_call_checks" - ) as mock_pre_call_checks: - with patch.object( - router, "_get_client", return_value=None - ) as mock_get_client: + with patch.object(router, "_update_kwargs_with_deployment") as mock_update_kwargs: + with patch.object(router, "async_routing_strategy_pre_call_checks") as mock_pre_call_checks: + with patch.object(router, "_get_client", return_value=None) as mock_get_client: result = await router._ageneric_api_call_with_fallbacks_helper( model="gpt-3.5-turbo", original_generic_function=mock_generic_function, @@ -1392,7 +1331,7 @@ async def test_router_ageneric_api_call_with_fallbacks_helper(): with patch.object(router, "async_get_available_deployment") as mock_get_deployment: mock_get_deployment.side_effect = Exception("No deployment available") - with pytest.raises(Exception, match='No deployment available') as exc_info: + with pytest.raises(Exception, match="No deployment available") as exc_info: await router._ageneric_api_call_with_fallbacks_helper( model="gpt-3.5-turbo", original_generic_function=mock_generic_function, @@ -1420,15 +1359,9 @@ async def test_router_ageneric_api_call_with_fallbacks_helper(): mock_semaphore = asyncio.Semaphore(1) - with patch.object( - router, "_update_kwargs_with_deployment" - ) as mock_update_kwargs: - with patch.object( - router, "_get_client", return_value=mock_semaphore - ) as mock_get_client: - with patch.object( - router, "async_routing_strategy_pre_call_checks" - ) as mock_pre_call_checks: + with patch.object(router, "_update_kwargs_with_deployment") as mock_update_kwargs: + with patch.object(router, "_get_client", return_value=mock_semaphore) as mock_get_client: + with patch.object(router, "async_routing_strategy_pre_call_checks") as mock_pre_call_checks: result = await router._ageneric_api_call_with_fallbacks_helper( model="gpt-3.5-turbo", original_generic_function=mock_semaphore_function, @@ -1457,16 +1390,10 @@ async def test_router_ageneric_api_call_with_fallbacks_helper(): }, } - with patch.object( - router, "_update_kwargs_with_deployment" - ) as mock_update_kwargs: - with patch.object( - router, "_get_client", return_value=None - ) as mock_get_client: - with patch.object( - router, "async_routing_strategy_pre_call_checks" - ) as mock_pre_call_checks: - with pytest.raises(Exception, match='Mock failure') as exc_info: + with patch.object(router, "_update_kwargs_with_deployment") as mock_update_kwargs: + with patch.object(router, "_get_client", return_value=None) as mock_get_client: + with patch.object(router, "async_routing_strategy_pre_call_checks") as mock_pre_call_checks: + with pytest.raises(Exception, match="Mock failure") as exc_info: await router._ageneric_api_call_with_fallbacks_helper( model="gpt-3.5-turbo", original_generic_function=mock_failing_function, @@ -1534,9 +1461,9 @@ async def test_ageneric_api_call_deployment_model_overrides_alias(): original_generic_function=capture_model, ) - assert ( - captured["model"] == "vertex_ai/gemini-2.5-flash" - ), f"Expected deployment model 'vertex_ai/gemini-2.5-flash', got '{captured['model']}'" + assert captured["model"] == "vertex_ai/gemini-2.5-flash", ( + f"Expected deployment model 'vertex_ai/gemini-2.5-flash', got '{captured['model']}'" + ) @pytest.mark.asyncio @@ -1642,14 +1569,10 @@ def test_router_get_model_access_groups_team_only_models(): ] ) - access_groups = router.get_model_access_groups( - model_name="gpt-3.5-turbo", team_id=None - ) + access_groups = router.get_model_access_groups(model_name="gpt-3.5-turbo", team_id=None) assert len(access_groups) == 0 - access_groups = router.get_model_access_groups( - model_name="gpt-3.5-turbo", team_id="team_1" - ) + access_groups = router.get_model_access_groups(model_name="gpt-3.5-turbo", team_id="team_1") assert list(access_groups.keys()) == ["default-models"] @@ -1744,9 +1667,7 @@ def test_model_group_info_cost_from_db_model_info(): ] ) - with patch.object( - router, "get_deployment_model_info", side_effect=Exception("not found") - ): + with patch.object(router, "get_deployment_model_info", side_effect=Exception("not found")): result = router._cached_get_model_group_info("my-custom-model") assert result is not None assert result.input_cost_per_token == 0.0001 @@ -1774,9 +1695,7 @@ def test_model_group_info_cost_none_when_db_model_info_has_no_cost(): ] ) - with patch.object( - router, "get_deployment_model_info", side_effect=Exception("not found") - ): + with patch.object(router, "get_deployment_model_info", side_effect=Exception("not found")): result = router._cached_get_model_group_info("my-custom-model-no-cost") assert result is not None assert result.input_cost_per_token is None @@ -1846,9 +1765,7 @@ def test_model_group_info_with_stringified_cost_values(): } return None - with patch.object( - router, "get_deployment_model_info", side_effect=_model_info_with_str_costs - ): + with patch.object(router, "get_deployment_model_info", side_effect=_model_info_with_str_costs): result = router._set_model_group_info( model_group="my-custom-model", user_facing_model_group_name="my-custom-model", @@ -1894,9 +1811,7 @@ def test_model_group_info_db_fallback_with_stringified_cost_values(): ] ) - with patch.object( - router, "get_deployment_model_info", side_effect=Exception("not found") - ): + with patch.object(router, "get_deployment_model_info", side_effect=Exception("not found")): result = router._set_model_group_info( model_group="my-custom-model", user_facing_model_group_name="my-custom-model", @@ -2156,6 +2071,7 @@ async def test_acompletion_streaming_iterator(): # Collect streamed chunks — the first chunk succeeds, then the error re-raises collected_chunks = [] + async def _drain(): async for chunk in result: collected_chunks.append(chunk) @@ -2849,11 +2765,7 @@ def _make_responses_iterator( BaseResponsesAPIStreamingIterator, ) - base = ( - LiteLLMCompletionStreamingIterator - if bridge - else BaseResponsesAPIStreamingIterator - ) + base = LiteLLMCompletionStreamingIterator if bridge else BaseResponsesAPIStreamingIterator class _Iter(base): def __init__(self): @@ -2933,9 +2845,7 @@ async def test_aresponses_streaming_iterator_fallback(): BaseResponsesAPIStreamingIterator, ) - router = _make_router_with_fallback( - "anthropic/claude-sonnet-4-6", "vertex_ai/claude-sonnet-4-6" - ) + router = _make_router_with_fallback("anthropic/claude-sonnet-4-6", "vertex_ai/claude-sonnet-4-6") src = _make_responses_iterator( chunks=[MagicMock(type="response.created")], error=MidStreamFallbackError( @@ -3018,9 +2928,9 @@ async def test_aresponses_streaming_iterator_writes_litellm_metadata_on_fallback fbk = mock_fallback_utils.call_args.kwargs["kwargs"] assert "litellm_metadata" in fbk, "wrong metadata_variable_name" assert fbk["litellm_metadata"]["model_group"] == "gpt-4" - assert "model_group" not in fbk.get( - "metadata", {} - ), "model_group leaked into 'metadata' instead of 'litellm_metadata'" + assert "model_group" not in fbk.get("metadata", {}), ( + "model_group leaked into 'metadata' instead of 'litellm_metadata'" + ) @pytest.mark.asyncio @@ -3138,9 +3048,7 @@ async def test_aresponses_streaming_iterator_combines_partial_usage(): fallback_response_object = ResponsesAPIResponse( id="resp_test", created_at=0, model="gpt-4", object="response", output=[] ) - fallback_response_object.usage = ResponseAPIUsage( - input_tokens=20, output_tokens=15, total_tokens=35 - ) + fallback_response_object.usage = ResponseAPIUsage(input_tokens=20, output_tokens=15, total_tokens=35) fallback_event = ResponseCompletedEvent( type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, response=fallback_response_object, @@ -3149,9 +3057,7 @@ async def test_aresponses_streaming_iterator_combines_partial_usage(): with ( patch( "litellm.main.stream_chunk_builder", - return_value=SimpleNamespace( - usage=SimpleNamespace(prompt_tokens=10, completion_tokens=4) - ), + return_value=SimpleNamespace(usage=SimpleNamespace(prompt_tokens=10, completion_tokens=4)), ), patch.object( router, @@ -3452,9 +3358,7 @@ def test_pre_call_checks_skips_token_count_without_max_input_tokens(monkeypatch) monkeypatch.setattr(router, "get_router_model_info", lambda **kwargs: {}) calls = [] - monkeypatch.setattr( - litellm, "token_counter", lambda *a, **k: calls.append(1) or 1000 - ) + monkeypatch.setattr(litellm, "token_counter", lambda *a, **k: calls.append(1) or 1000) deployments = [ {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, @@ -3482,14 +3386,10 @@ def test_pre_call_checks_counts_once_and_filters_on_max_input_tokens(monkeypatch ], enable_pre_call_checks=True, ) - monkeypatch.setattr( - router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 5} - ) + monkeypatch.setattr(router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 5}) calls = [] - monkeypatch.setattr( - litellm, "token_counter", lambda *a, **k: calls.append(1) or 1000 - ) + monkeypatch.setattr(litellm, "token_counter", lambda *a, **k: calls.append(1) or 1000) deployments = [ {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, @@ -3516,14 +3416,10 @@ def test_pre_call_checks_uses_precounted_tokens(monkeypatch): ], enable_pre_call_checks=True, ) - monkeypatch.setattr( - router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 5} - ) + monkeypatch.setattr(router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 5}) calls = [] - monkeypatch.setattr( - litellm, "token_counter", lambda *a, **k: calls.append(1) or 1 - ) + monkeypatch.setattr(litellm, "token_counter", lambda *a, **k: calls.append(1) or 1) deployments = [ {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, @@ -3550,9 +3446,7 @@ async def test_async_get_healthy_deployments_counts_tokens_off_the_event_loop(mo ], enable_pre_call_checks=True, ) - monkeypatch.setattr( - router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 1_000_000} - ) + monkeypatch.setattr(router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 1_000_000}) counting_threads = [] monkeypatch.setattr( @@ -3648,14 +3542,10 @@ def test_pre_call_checks_does_not_recount_inline_after_an_off_loop_failure(monke ], enable_pre_call_checks=True, ) - monkeypatch.setattr( - router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 5} - ) + monkeypatch.setattr(router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 5}) calls = [] - monkeypatch.setattr( - litellm, "token_counter", lambda *a, **k: calls.append(1) or 1000 - ) + monkeypatch.setattr(litellm, "token_counter", lambda *a, **k: calls.append(1) or 1000) deployments = [ {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, @@ -3683,9 +3573,7 @@ async def test_async_get_healthy_deployments_never_recounts_on_the_loop(monkeypa ], enable_pre_call_checks=True, ) - monkeypatch.setattr( - router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 5} - ) + monkeypatch.setattr(router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 5}) counting_threads = [] @@ -3720,9 +3608,7 @@ async def test_acount_pre_call_check_tokens_leaves_the_event_loop_free(monkeypat ], enable_pre_call_checks=True, ) - monkeypatch.setattr( - router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 5} - ) + monkeypatch.setattr(router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 5}) deployments = [ {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, @@ -3758,9 +3644,7 @@ async def test_acount_pre_call_check_tokens_skips_without_max_input_tokens(monke monkeypatch.setattr(router, "get_router_model_info", lambda **kwargs: {}) calls = [] - monkeypatch.setattr( - litellm, "token_counter", lambda *a, **k: calls.append(1) or 1000 - ) + monkeypatch.setattr(litellm, "token_counter", lambda *a, **k: calls.append(1) or 1000) count = await router._acount_pre_call_check_tokens( model="m", @@ -3788,9 +3672,7 @@ def test_pre_call_checks_counts_tokens_from_responses_input_string(monkeypatch): ], enable_pre_call_checks=True, ) - monkeypatch.setattr( - router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 1} - ) + monkeypatch.setattr(router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 1}) deployments = [ {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, @@ -3815,9 +3697,7 @@ def test_pre_call_checks_counts_tokens_from_responses_input_list(monkeypatch): ], enable_pre_call_checks=True, ) - monkeypatch.setattr( - router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 1} - ) + monkeypatch.setattr(router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 1}) deployments = [ {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, @@ -3859,9 +3739,7 @@ def test_pre_call_checks_counts_responses_instructions_tokens(monkeypatch): ) assert with_instructions_tokens > input_only_tokens - monkeypatch.setattr( - router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": input_only_tokens} - ) + monkeypatch.setattr(router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": input_only_tokens}) with pytest.raises(litellm.ContextWindowExceededError): router._pre_call_checks( model="m", @@ -3930,9 +3808,7 @@ def test_pre_call_checks_counts_tool_definition_tokens(monkeypatch, prompt_kwarg prompt_only_tokens = router._count_pre_call_check_tokens( messages=prompt_kwargs.get("messages"), input=prompt_kwargs.get("input") ) - monkeypatch.setattr( - router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": prompt_only_tokens} - ) + monkeypatch.setattr(router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": prompt_only_tokens}) assert len(router._pre_call_checks(model="m", healthy_deployments=deployments, **prompt_kwargs)) == 1 with pytest.raises(litellm.ContextWindowExceededError): @@ -4052,7 +3928,7 @@ def test_count_pre_call_check_tokens_across_api_surfaces(): assert string_input_tokens > 0 assert list_input_tokens > 0 - with pytest.raises(ValueError, match='Either messages or input must be provided to count tokens'): + with pytest.raises(ValueError, match="Either messages or input must be provided to count tokens"): router._count_pre_call_check_tokens(messages=None, input=None) @@ -4067,9 +3943,7 @@ def test_pre_call_checks_no_messages_or_input_does_not_crash(monkeypatch): ], enable_pre_call_checks=True, ) - monkeypatch.setattr( - router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 5} - ) + monkeypatch.setattr(router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 5}) counted: list[dict] = [] original = router._count_pre_call_check_tokens @@ -4159,9 +4033,7 @@ def test_get_deployment_model_info_base_model_flow(): } # Test Case 1: Base model flow with custom model info that has base_model - with patch.object( - litellm, "model_cost", {"test-custom-model": mock_custom_model_info} - ): + with patch.object(litellm, "model_cost", {"test-custom-model": mock_custom_model_info}): with patch.object(litellm, "get_model_info") as mock_get_model_info: # Configure mock returns mock_get_model_info.side_effect = lambda model: { @@ -4169,15 +4041,11 @@ def test_get_deployment_model_info_base_model_flow(): "test-model": mock_litellm_model_name_info, }.get(model) - result = router.get_deployment_model_info( - model_id="test-custom-model", model_name="test-model" - ) + result = router.get_deployment_model_info(model_id="test-custom-model", model_name="test-model") # Verify that get_model_info was called for both base model and model name assert mock_get_model_info.call_count == 2 - mock_get_model_info.assert_any_call( - model="gpt-3.5-turbo" - ) # base model call + mock_get_model_info.assert_any_call(model="gpt-3.5-turbo") # base model call mock_get_model_info.assert_any_call(model="test-model") # model name call # Verify the result contains merged information @@ -4188,26 +4056,18 @@ def test_get_deployment_model_info_base_model_flow(): # 2. The result of step 1 gets merged into litellm_model_name_info (custom+base override litellm) # Fields from custom model (should override base model values) - assert ( - result["input_cost_per_token"] == 0.001 - ) # From custom model (overrides base 0.0015) - assert ( - result["output_cost_per_token"] == 0.002 - ) # From custom model (same as base) + assert result["input_cost_per_token"] == 0.001 # From custom model (overrides base 0.0015) + assert result["output_cost_per_token"] == 0.002 # From custom model (same as base) assert result["custom_field"] == "custom_value" # From custom model # Fields from base model that weren't overridden by custom assert result["max_tokens"] == 4096 # From base model assert result["litellm_provider"] == "openai" # From base model - assert ( - result["mode"] == "chat" - ) # From base model (overrides litellm "completion") + assert result["mode"] == "chat" # From base model (overrides litellm "completion") # The key field comes from base model since both base and litellm have it # and base model info overrides litellm model name info in final merge - assert ( - result["key"] == "gpt-3.5-turbo" - ) # From base model (overrides litellm key) + assert result["key"] == "gpt-3.5-turbo" # From base model (overrides litellm key) # Test Case 2: Custom model info without base_model mock_custom_model_info_no_base = { @@ -4226,9 +4086,7 @@ def test_get_deployment_model_info_base_model_flow(): "test-model": mock_litellm_model_name_info, }.get(model) - result = router.get_deployment_model_info( - model_id="test-custom-model-no-base", model_name="test-model" - ) + result = router.get_deployment_model_info(model_id="test-custom-model-no-base", model_name="test-model") # Should only call get_model_info once for model name (no base model) assert mock_get_model_info.call_count == 1 @@ -4248,9 +4106,7 @@ def test_get_deployment_model_info_base_model_flow(): "test-model": mock_litellm_model_name_info, }.get(model) - result = router.get_deployment_model_info( - model_id="non-existent-model", model_name="test-model" - ) + result = router.get_deployment_model_info(model_id="non-existent-model", model_name="test-model") # Should only call get_model_info once for model name assert mock_get_model_info.call_count == 1 @@ -4283,9 +4139,7 @@ def test_get_deployment_model_info_base_model_flow(): mock_get_model_info.side_effect = mock_get_model_info_side_effect - result = router.get_deployment_model_info( - model_id="test-custom-model-invalid", model_name="test-model" - ) + result = router.get_deployment_model_info(model_id="test-custom-model-invalid", model_name="test-model") # Should handle exception gracefully and still return merged result assert result is not None @@ -4294,12 +4148,8 @@ def test_get_deployment_model_info_base_model_flow(): # Test Case 5: Both model_cost.get() and get_model_info() return None with patch.object(litellm, "model_cost", {}): - with patch.object( - litellm, "get_model_info", side_effect=Exception("Not found") - ): - result = router.get_deployment_model_info( - model_id="non-existent", model_name="non-existent" - ) + with patch.object(litellm, "get_model_info", side_effect=Exception("Not found")): + result = router.get_deployment_model_info(model_id="non-existent", model_name="non-existent") # Should return None when no model info is found assert result is None @@ -4322,9 +4172,7 @@ def test_get_deployment_model_info_base_model_flow(): # Model NOT in built-in cost map — raise exception mock_get_model_info.side_effect = Exception("Model not in cost map") - result = router.get_deployment_model_info( - model_id="custom-model-id", model_name="unknown-model" - ) + result = router.get_deployment_model_info(model_id="custom-model-id", model_name="unknown-model") # Should return custom_model_info even when litellm_model_name_model_info is None assert result is not None @@ -4360,15 +4208,11 @@ def test_get_deployment_model_info_base_model_flow(): mock_get_model_info.side_effect = get_info_side_effect - result = router.get_deployment_model_info( - model_id="custom-with-base", model_name="unknown-model" - ) + result = router.get_deployment_model_info(model_id="custom-with-base", model_name="unknown-model") # Should return custom_model_info merged with base model info assert result is not None - assert ( - result["input_cost_per_token"] == 0.01 - ) # From custom (overrides base) + assert result["input_cost_per_token"] == 0.01 # From custom (overrides base) assert result["max_tokens"] == 8192 # From base model assert result["litellm_provider"] == "openai" # From base model @@ -4415,18 +4259,14 @@ def test_get_deployment_model_info_base_model_merge_priority(): "litellm_only_field": "litellm_value", } - with patch.object( - litellm, "model_cost", {"custom-model-id": mock_custom_model_info} - ): + with patch.object(litellm, "model_cost", {"custom-model-id": mock_custom_model_info}): with patch.object(litellm, "get_model_info") as mock_get_model_info: mock_get_model_info.side_effect = lambda model: { "gpt-4": mock_base_model_info, "test-model": mock_litellm_model_name_info, }.get(model) - result = router.get_deployment_model_info( - model_id="custom-model-id", model_name="test-model" - ) + result = router.get_deployment_model_info(model_id="custom-model-id", model_name="test-model") assert result is not None @@ -4436,29 +4276,17 @@ def test_get_deployment_model_info_base_model_merge_priority(): # 3. Result from steps 1-2 overrides litellm_model_name_info # Fields that should come from custom model info (highest priority) - assert ( - result["input_cost_per_token"] == 0.01 - ) # From custom model (overrides base 0.03) - assert ( - result["max_tokens"] == 8000 - ) # From custom model (overrides base 4096) + assert result["input_cost_per_token"] == 0.01 # From custom model (overrides base 0.03) + assert result["max_tokens"] == 8000 # From custom model (overrides base 4096) assert result["custom_only_field"] == "custom_value" # From custom model # Fields that should come from base model (not overridden by custom) - assert ( - result["output_cost_per_token"] == 0.06 - ) # From base model (not in custom) - assert ( - result["litellm_provider"] == "openai" - ) # From base model (not in custom) - assert ( - result["base_only_field"] == "base_value" - ) # From base model (not in custom) + assert result["output_cost_per_token"] == 0.06 # From base model (not in custom) + assert result["litellm_provider"] == "openai" # From base model (not in custom) + assert result["base_only_field"] == "base_value" # From base model (not in custom) # Fields that should come from litellm model name info (not overridden by custom+base) - assert ( - result["mode"] == "completion" - ) # From litellm model name info (not in custom or base) + assert result["mode"] == "completion" # From litellm model name info (not in custom or base) assert ( result["litellm_only_field"] == "litellm_value" ) # From litellm model name info (not in custom or base) @@ -4495,10 +4323,9 @@ def test_add_deployment_model_to_endpoint_for_llm_passthrough_route(): model="special-bedrock-model", model_name="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", ) - assert ( - result["endpoint"] - == "/model/us.anthropic.claude-haiku-4-5-20251001-v1:0/invoke" - ), f"Expected '/model/us.anthropic.claude-haiku-4-5-20251001-v1:0/invoke', got '{result['endpoint']}'" + assert result["endpoint"] == "/model/us.anthropic.claude-haiku-4-5-20251001-v1:0/invoke", ( + f"Expected '/model/us.anthropic.claude-haiku-4-5-20251001-v1:0/invoke', got '{result['endpoint']}'" + ) # Test Case 2: Bedrock invoke-with-response-stream endpoint kwargs = { @@ -4510,10 +4337,9 @@ def test_add_deployment_model_to_endpoint_for_llm_passthrough_route(): model="special-bedrock-model", model_name="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", ) - assert ( - result["endpoint"] - == "/model/us.anthropic.claude-haiku-4-5-20251001-v1:0/invoke-with-response-stream" - ), f"Expected streaming endpoint with stripped prefix, got '{result['endpoint']}'" + assert result["endpoint"] == "/model/us.anthropic.claude-haiku-4-5-20251001-v1:0/invoke-with-response-stream", ( + f"Expected streaming endpoint with stripped prefix, got '{result['endpoint']}'" + ) # Test Case 3: Bedrock converse endpoint kwargs = { @@ -4525,9 +4351,9 @@ def test_add_deployment_model_to_endpoint_for_llm_passthrough_route(): model="bedrock-model", model_name="bedrock/us.meta.llama3-8b-instruct-v1:0", ) - assert ( - result["endpoint"] == "/model/us.meta.llama3-8b-instruct-v1:0/converse" - ), f"Expected '/model/us.meta.llama3-8b-instruct-v1:0/converse', got '{result['endpoint']}'" + assert result["endpoint"] == "/model/us.meta.llama3-8b-instruct-v1:0/converse", ( + f"Expected '/model/us.meta.llama3-8b-instruct-v1:0/converse', got '{result['endpoint']}'" + ) # Test Case 4: Bedrock provider prefix auto-detected from model_name kwargs = { @@ -4538,9 +4364,9 @@ def test_add_deployment_model_to_endpoint_for_llm_passthrough_route(): model="router-model", model_name="bedrock/us.meta.llama3-8b-instruct-v1:0", ) - assert ( - result["endpoint"] == "/model/us.meta.llama3-8b-instruct-v1:0/invoke" - ), f"Expected '/model/us.meta.llama3-8b-instruct-v1:0/invoke', got '{result['endpoint']}'" + assert result["endpoint"] == "/model/us.meta.llama3-8b-instruct-v1:0/invoke", ( + f"Expected '/model/us.meta.llama3-8b-instruct-v1:0/invoke', got '{result['endpoint']}'" + ) def test_update_kwargs_with_deployment_uses_pass_through_request_timeout(): @@ -4592,14 +4418,10 @@ async def test_router_acompletion_with_unknown_model_and_default_fallback(): # Initialize the router with a default fallback router = litellm.Router(model_list=model_list, default_fallbacks=["gpt-4o"]) - messages = [ - {"role": "user", "content": "This call should succeed by falling back."} - ] + messages = [{"role": "user", "content": "This call should succeed by falling back."}] # Call completion with a model name that is NOT in the model_list - response = await router.acompletion( - model="completely-unknown-model", messages=messages - ) + response = await router.acompletion(model="completely-unknown-model", messages=messages) # Check that the call did not fail and we received a valid response object. assert response is not None @@ -4691,15 +4513,10 @@ def test_get_deployment_credentials_with_provider_aws_bedrock_runtime_endpoint() ], ) - credentials = router.get_deployment_credentials_with_provider( - model_id="bedrock-claude-model" - ) + credentials = router.get_deployment_credentials_with_provider(model_id="bedrock-claude-model") assert credentials is not None - assert ( - credentials["aws_bedrock_runtime_endpoint"] - == "https://bedrock-runtime.us-east-1.amazonaws.com" - ) + assert credentials["aws_bedrock_runtime_endpoint"] == "https://bedrock-runtime.us-east-1.amazonaws.com" assert credentials["aws_access_key_id"] == "test-access-key" assert credentials["aws_secret_access_key"] == "test-secret-key" assert credentials["aws_region_name"] == "us-east-1" @@ -4726,9 +4543,7 @@ def test_get_deployment_credentials_with_provider_includes_bucket_name(): ], ) - credentials = router.get_deployment_credentials_with_provider( - model_id="vertex-gemini" - ) + credentials = router.get_deployment_credentials_with_provider(model_id="vertex-gemini") assert credentials is not None assert credentials["gcs_bucket_name"] == "my-batch-bucket" @@ -4768,9 +4583,7 @@ def test_get_deployment_credentials_with_provider_resolves_credential_name(): ], ) - credentials = router.get_deployment_credentials_with_provider( - model_id="azure-gpt-4" - ) + credentials = router.get_deployment_credentials_with_provider(model_id="azure-gpt-4") assert credentials is not None assert credentials["api_key"] == "resolved-api-key" @@ -4806,9 +4619,7 @@ def test_get_deployment_credentials_with_provider_bedrock_batch_fields(): ], ) - credentials = router.get_deployment_credentials_with_provider( - model_id="bedrock-batch-model" - ) + credentials = router.get_deployment_credentials_with_provider(model_id="bedrock-batch-model") assert credentials is not None assert credentials["custom_llm_provider"] == "bedrock" @@ -4851,9 +4662,7 @@ def test_get_deployment_credentials_with_provider_preserves_aws_auth_params(): ], ) - credentials = router.get_deployment_credentials_with_provider( - model_id="bedrock-batch-model" - ) + credentials = router.get_deployment_credentials_with_provider(model_id="bedrock-batch-model") assert credentials is not None for key, value in aws_auth_params.items(): @@ -4888,15 +4697,11 @@ def test_get_deployment_credentials_with_provider_team_wildcard_priority(): ], ) - team_credentials = router.get_deployment_credentials_with_provider( - model_id="openai/gpt-5.2", team_id="team-1" - ) + team_credentials = router.get_deployment_credentials_with_provider(model_id="openai/gpt-5.2", team_id="team-1") assert team_credentials is not None assert team_credentials["api_key"] == "team-key" - global_credentials = router.get_deployment_credentials_with_provider( - model_id="openai/gpt-5.2" - ) + global_credentials = router.get_deployment_credentials_with_provider(model_id="openai/gpt-5.2") assert global_credentials is not None assert global_credentials["api_key"] == "global-key" @@ -4937,15 +4742,11 @@ def test_get_deployment_credentials_with_provider_skips_other_team_deployment(): assert other_team_credentials is not None assert other_team_credentials["vertex_project"] == "shared-project" - unscoped_credentials = router.get_deployment_credentials_with_provider( - model_id="gemini-2.5-pro" - ) + unscoped_credentials = router.get_deployment_credentials_with_provider(model_id="gemini-2.5-pro") assert unscoped_credentials is not None assert unscoped_credentials["vertex_project"] == "shared-project" - owner_credentials = router.get_deployment_credentials_with_provider( - model_id="gemini-2.5-pro", team_id="team-b" - ) + owner_credentials = router.get_deployment_credentials_with_provider(model_id="gemini-2.5-pro", team_id="team-b") assert owner_credentials is not None assert owner_credentials["vertex_project"] == "team-b-project" @@ -4972,16 +4773,8 @@ def test_get_deployment_credentials_with_provider_no_fallback_to_other_team_only ], ) - assert ( - router.get_deployment_credentials_with_provider( - model_id="gemini-2.5-pro", team_id="team-a" - ) - is None - ) - assert ( - router.get_deployment_credentials_with_provider(model_id="gemini-2.5-pro") - is None - ) + assert router.get_deployment_credentials_with_provider(model_id="gemini-2.5-pro", team_id="team-a") is None + assert router.get_deployment_credentials_with_provider(model_id="gemini-2.5-pro") is None def test_deployment_usable_by_team_helpers(): @@ -5021,9 +4814,7 @@ def test_deployment_usable_by_team_helpers(): assert router._deployment_usable_by_team(shared, "team-a") is True assert router._deployment_usable_by_team(shared, None) is True - picked = router._get_model_group_deployment_usable_by_team( - model_group_name="gemini-2.5-pro", team_id="team-a" - ) + picked = router._get_model_group_deployment_usable_by_team(model_group_name="gemini-2.5-pro", team_id="team-a") assert picked is not None assert picked.litellm_params.vertex_project == "shared-project" @@ -5033,12 +4824,7 @@ def test_deployment_usable_by_team_helpers(): assert owner_picked is not None assert owner_picked.litellm_params.vertex_project == "team-b-project" - assert ( - router._get_model_group_deployment_usable_by_team( - model_group_name="unknown-model", team_id="team-a" - ) - is None - ) + assert router._get_model_group_deployment_usable_by_team(model_group_name="unknown-model", team_id="team-a") is None def test_get_deployment_credentials_with_provider_skips_other_team_wildcard(): @@ -5070,9 +4856,7 @@ def test_get_deployment_credentials_with_provider_skips_other_team_wildcard(): assert other_team_credentials is not None assert other_team_credentials["api_key"] == "global-key" - owner_credentials = router.get_deployment_credentials_with_provider( - model_id="openai/gpt-5.2", team_id="team-b" - ) + owner_credentials = router.get_deployment_credentials_with_provider(model_id="openai/gpt-5.2", team_id="team-b") assert owner_credentials is not None assert owner_credentials["api_key"] == "team-b-key" @@ -5084,21 +4868,11 @@ def test_team_wildcard_credentials_not_usable_after_delete_deployment(): """ router = litellm.Router(model_list=[_team_wildcard_model(api_key="old-key")]) - assert ( - router.get_deployment_credentials_with_provider( - model_id="openai/gpt-5.2", team_id="team-1" - ) - is not None - ) + assert router.get_deployment_credentials_with_provider(model_id="openai/gpt-5.2", team_id="team-1") is not None router.delete_deployment(id="team-wildcard-id") - assert ( - router.get_deployment_credentials_with_provider( - model_id="openai/gpt-5.2", team_id="team-1" - ) - is None - ) + assert router.get_deployment_credentials_with_provider(model_id="openai/gpt-5.2", team_id="team-1") is None def test_global_wildcard_pattern_router_evicts_stale_entry_on_upsert_and_delete(): @@ -5171,22 +4945,13 @@ def test_team_wildcard_credentials_refreshed_on_upsert_and_set_model_list(): router = litellm.Router(model_list=[_team_wildcard_model(api_key="old-key")]) - router.upsert_deployment( - deployment=Deployment(**_team_wildcard_model(api_key="new-key")) - ) - credentials = router.get_deployment_credentials_with_provider( - model_id="openai/gpt-5.2", team_id="team-1" - ) + router.upsert_deployment(deployment=Deployment(**_team_wildcard_model(api_key="new-key"))) + credentials = router.get_deployment_credentials_with_provider(model_id="openai/gpt-5.2", team_id="team-1") assert credentials is not None assert credentials["api_key"] == "new-key" router.set_model_list(model_list=[]) - assert ( - router.get_deployment_credentials_with_provider( - model_id="openai/gpt-5.2", team_id="team-1" - ) - is None - ) + assert router.get_deployment_credentials_with_provider(model_id="openai/gpt-5.2", team_id="team-1") is None def test_get_available_guardrail_single_deployment(): @@ -5375,9 +5140,7 @@ async def test_anthropic_messages_call_type_is_cached(): startTime=1234567890.0, endTime=1234567891.0, completionStartTime=1234567890.5, - model_map_information=StandardLoggingModelInformation( - model_map_key="gpt-3.5-turbo", model_map_value=None - ), + model_map_information=StandardLoggingModelInformation(model_map_key="gpt-3.5-turbo", model_map_value=None), model="gpt-3.5-turbo", model_id="model-123", model_group="openai-gpt", @@ -5456,12 +5219,8 @@ async def test_anthropic_messages_call_type_is_cached(): ) # This assertion will FAIL if anthropic_messages is filtered out - assert ( - cached_result is not None - ), "Model ID should be cached for anthropic_messages call type" - assert ( - cached_result["model_id"] == test_model_id - ), f"Expected {test_model_id}, got {cached_result['model_id']}" + assert cached_result is not None, "Model ID should be cached for anthropic_messages call type" + assert cached_result["model_id"] == test_model_id, f"Expected {test_model_id}, got {cached_result['model_id']}" def test_update_kwargs_with_deployment_propagates_model_tags(): @@ -5486,9 +5245,7 @@ def test_update_kwargs_with_deployment_propagates_model_tags(): ) kwargs: dict = {"metadata": {}} - deployment = router.get_deployment_by_model_group_name( - model_group_name="gpt-4o-mini" - ) + deployment = router.get_deployment_by_model_group_name(model_group_name="gpt-4o-mini") router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) # Deployment tags should be propagated to kwargs metadata @@ -5517,9 +5274,7 @@ def test_update_kwargs_with_deployment_merges_tags_without_duplicates(): # Simulate request that already has tags (from request body or key/team level) kwargs: dict = {"metadata": {"tags": ["user-tag", "shared-tag"]}} - deployment = router.get_deployment_by_model_group_name( - model_group_name="gpt-4o-mini" - ) + deployment = router.get_deployment_by_model_group_name(model_group_name="gpt-4o-mini") router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) # Both sources should be merged, no duplicates @@ -5546,9 +5301,7 @@ def test_update_kwargs_with_deployment_no_tags(): ) kwargs: dict = {"metadata": {}} - deployment = router.get_deployment_by_model_group_name( - model_group_name="gpt-4o-mini" - ) + deployment = router.get_deployment_by_model_group_name(model_group_name="gpt-4o-mini") router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) # No tags key should be added if deployment has no tags @@ -5586,9 +5339,7 @@ def test_update_kwargs_with_deployment_merges_tools(): }, ], } - deployment = router.get_deployment_by_model_group_name( - model_group_name="o3-deep-research" - ) + deployment = router.get_deployment_by_model_group_name(model_group_name="o3-deep-research") router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) # Tools should be merged: deployment first, then request @@ -5619,9 +5370,7 @@ def test_update_kwargs_with_deployment_merge_tools_deployment_only(): ) kwargs: dict = {"metadata": {}} - deployment = router.get_deployment_by_model_group_name( - model_group_name="o3-deep-research" - ) + deployment = router.get_deployment_by_model_group_name(model_group_name="o3-deep-research") router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) assert kwargs["tools"] == [{"type": "web_search"}] @@ -5650,9 +5399,7 @@ def test_update_kwargs_with_deployment_merge_tools_request_overrides_tool_choice "metadata": {}, "tool_choice": "none", } - deployment = router.get_deployment_by_model_group_name( - model_group_name="o3-deep-research" - ) + deployment = router.get_deployment_by_model_group_name(model_group_name="o3-deep-research") router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) # Request tool_choice should be preserved (merged tools still applied) @@ -5754,12 +5501,8 @@ def test_update_kwargs_with_deployment_model_info_in_litellm_metadata(): ) kwargs: dict = {} - deployment = router.get_deployment_by_model_group_name( - model_group_name="claude-sonnet-4" - ) - router._update_kwargs_with_deployment( - deployment=deployment, kwargs=kwargs, function_name="generic_api_call" - ) + deployment = router.get_deployment_by_model_group_name(model_group_name="claude-sonnet-4") + router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs, function_name="generic_api_call") assert "litellm_metadata" in kwargs model_info = kwargs["litellm_metadata"]["model_info"] @@ -5791,12 +5534,8 @@ def test_update_kwargs_with_deployment_model_info_in_metadata(): ) kwargs: dict = {} - deployment = router.get_deployment_by_model_group_name( - model_group_name="claude-sonnet-4" - ) - router._update_kwargs_with_deployment( - deployment=deployment, kwargs=kwargs, function_name=None - ) + deployment = router.get_deployment_by_model_group_name(model_group_name="claude-sonnet-4") + router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs, function_name=None) assert "metadata" in kwargs model_info = kwargs["metadata"]["model_info"] @@ -5909,6 +5648,7 @@ async def test_acompletion_streaming_iterator_does_not_log_success_on_terminal_f initial_kwargs=dict(initial_kwargs), ) collected = [] + async def _drain(): async for chunk in result: collected.append(chunk) @@ -5935,6 +5675,7 @@ async def test_acompletion_streaming_iterator_does_not_log_success_on_terminal_f initial_kwargs=dict(initial_kwargs), ) collected = [] + async def _drain(): async for chunk in result: collected.append(chunk) @@ -6151,23 +5892,17 @@ def test_multiregion_team_deployments_unique_model_names(): assert len(deployments) == 0 # With team_id: O(n) scan finds BOTH regional deployments - deployments = router._get_all_deployments( - model_name="claude-sonnet", team_id="metis-team" - ) + deployments = router._get_all_deployments(model_name="claude-sonnet", team_id="metis-team") assert len(deployments) == 2 deployment_names = {d["model_name"] for d in deployments} assert deployment_names == {"metis-claude-us-east-1", "metis-claude-us-west-2"} # Each deployment has a unique ID (critical for cooldown/retry to work) deployment_ids = {d["model_info"]["id"] for d in deployments} - assert ( - len(deployment_ids) == 2 - ), "Each deployment must have a unique ID for cooldown tracking" + assert len(deployment_ids) == 2, "Each deployment must have a unique ID for cooldown tracking" # Wrong team: returns nothing - deployments = router._get_all_deployments( - model_name="claude-sonnet", team_id="other-team" - ) + deployments = router._get_all_deployments(model_name="claude-sonnet", team_id="other-team") assert len(deployments) == 0 @@ -6212,12 +5947,8 @@ async def test_multiregion_team_failover_between_regions(): ) # Verify the router finds both deployments for the team - deployments = router._get_all_deployments( - model_name="claude-sonnet", team_id="metis-team" - ) - assert ( - len(deployments) == 2 - ), "Router must find both regional deployments by team_public_model_name" + deployments = router._get_all_deployments(model_name="claude-sonnet", team_id="metis-team") + assert len(deployments) == 2, "Router must find both regional deployments by team_public_model_name" # Make a normal request — should succeed from one of the regions response = await router.acompletion( @@ -6342,9 +6073,7 @@ def test_explicit_model_access_does_not_force_access_group_filtering(): }, ) - deployment_groups = [ - d.get("model_info", {}).get("access_groups") for d in deployments - ] + deployment_groups = [d.get("model_info", {}).get("access_groups") for d in deployments] assert ["AG1"] in deployment_groups assert ["AG2"] in deployment_groups @@ -6389,9 +6118,7 @@ def test_access_group_filter_empty_does_not_bypass_via_litellm_model_fallback( orig_groups = router.get_model_access_groups - def fake_get_model_access_groups( - model_name=None, model_access_group=None, team_id=None - ): + def fake_get_model_access_groups(model_name=None, model_access_group=None, team_id=None): if model_name == "gpt-5" and model_access_group is None: return {"AG1": ["gpt-5"], "AG2": ["gpt-5"]} return orig_groups( @@ -6466,9 +6193,7 @@ def test_access_group_block_does_not_silently_use_default_fallback_model( orig_groups = router.get_model_access_groups - def fake_get_model_access_groups( - model_name=None, model_access_group=None, team_id=None - ): + def fake_get_model_access_groups(model_name=None, model_access_group=None, team_id=None): if model_name == "gpt-5" and model_access_group is None: return {"AG1": ["gpt-5"], "AG2": ["gpt-5"]} return orig_groups( @@ -6535,9 +6260,7 @@ def test_access_group_block_via_litellm_model_branch_does_not_use_default_fallba orig_groups = router.get_model_access_groups - def fake_get_model_access_groups( - model_name=None, model_access_group=None, team_id=None - ): + def fake_get_model_access_groups(model_name=None, model_access_group=None, team_id=None): if model_name == "gpt-5" and model_access_group is None: return {"AG1": ["gpt-5"], "AG2": ["gpt-5"]} return orig_groups( @@ -6592,9 +6315,7 @@ def test_try_early_resolve_deployments_for_model_not_in_names(): ) assert ( - router_in_names._try_early_resolve_deployments_for_model_not_in_names( - model="gpt-5", request_team_id=None - ) + router_in_names._try_early_resolve_deployments_for_model_not_in_names(model="gpt-5", request_team_id=None) is None ) assert ( @@ -6616,10 +6337,8 @@ def test_try_early_resolve_deployments_for_model_not_in_names(): ] ) - pattern_result = ( - pattern_router._try_early_resolve_deployments_for_model_not_in_names( - model="openai/gpt-4o-mini", request_team_id=None - ) + pattern_result = pattern_router._try_early_resolve_deployments_for_model_not_in_names( + model="openai/gpt-4o-mini", request_team_id=None ) assert pattern_result is not None resolved_model, pattern_deployments = pattern_result @@ -6645,10 +6364,8 @@ def test_try_early_resolve_deployments_for_model_not_in_names(): }, } - default_result = ( - default_router._try_early_resolve_deployments_for_model_not_in_names( - model="brand-new-model", request_team_id=None - ) + default_result = default_router._try_early_resolve_deployments_for_model_not_in_names( + model="brand-new-model", request_team_id=None ) assert default_result is not None resolved_model, default_deployment = default_result @@ -6656,10 +6373,7 @@ def test_try_early_resolve_deployments_for_model_not_in_names(): assert isinstance(default_deployment, dict) assert default_deployment["litellm_params"]["model"] == "brand-new-model" # The original default_deployment must not be mutated. - assert ( - default_router.default_deployment["litellm_params"]["model"] - == "openai/will-be-overridden" - ) + assert default_router.default_deployment["litellm_params"]["model"] == "openai/will-be-overridden" def _router_with_two_deployments(blocked_flags): @@ -6707,10 +6421,7 @@ def _seed_unhealthy_states(router, unhealthy_ids, timestamp=None): ts = timestamp if timestamp is not None else time.time() router.health_state_cache.set_deployment_health_states( - { - uid: {"is_healthy": False, "timestamp": ts, "reason": "test_unhealthy"} - for uid in unhealthy_ids - } + {uid: {"is_healthy": False, "timestamp": ts, "reason": "test_unhealthy"} for uid in unhealthy_ids} ) @@ -6781,9 +6492,7 @@ async def test_async_get_fully_unhealthy_model_names_noop_with_allowed_fails_pol @pytest.mark.asyncio async def test_async_get_healthy_deployments_skips_blocked_deployment(): router = _router_with_two_deployments([True, False]) - healthy, all_dep = await router._async_get_healthy_deployments( - model="gpt-4o", parent_otel_span=None - ) + healthy, all_dep = await router._async_get_healthy_deployments(model="gpt-4o", parent_otel_span=None) healthy_ids = [d["model_info"]["id"] for d in healthy] assert "dep-0" not in healthy_ids assert "dep-1" in healthy_ids @@ -6792,9 +6501,7 @@ async def test_async_get_healthy_deployments_skips_blocked_deployment(): def test_get_healthy_deployments_sync_skips_blocked_deployment(): router = _router_with_two_deployments([False, True]) - healthy, all_dep = router._get_healthy_deployments( - model="gpt-4o", parent_otel_span=None - ) + healthy, all_dep = router._get_healthy_deployments(model="gpt-4o", parent_otel_span=None) healthy_ids = [d["model_info"]["id"] for d in healthy] assert "dep-0" in healthy_ids assert "dep-1" not in healthy_ids @@ -6811,9 +6518,7 @@ def test_filter_blocked_deployments_drops_blocked_keeps_unblocked(): @pytest.mark.asyncio async def test_public_async_get_healthy_deployments_skips_blocked_on_primary_path(): router = _router_with_two_deployments([True, False]) - deployments = await router.async_get_healthy_deployments( - model="gpt-4o", request_kwargs={} - ) + deployments = await router.async_get_healthy_deployments(model="gpt-4o", request_kwargs={}) assert isinstance(deployments, list) ids = [d["model_info"]["id"] for d in deployments] assert "dep-0" not in ids @@ -6855,9 +6560,7 @@ def _router_with_two_pass_through_deployments(blocked_flags): def test_get_available_deployment_for_pass_through_skips_blocked(): router = _router_with_two_pass_through_deployments([True, False]) - deployment = router.get_available_deployment_for_pass_through( - model="gpt-4o", request_kwargs={} - ) + deployment = router.get_available_deployment_for_pass_through(model="gpt-4o", request_kwargs={}) assert deployment["model_info"]["id"] == "pt-1" @@ -6866,9 +6569,7 @@ def test_get_available_deployment_for_pass_through_raises_when_dict_blocked(): router = _router_with_two_pass_through_deployments([True, True]) with pytest.raises(litellm.ServiceUnavailableError): - router.get_available_deployment_for_pass_through( - model="pt-0", request_kwargs={} - ) + router.get_available_deployment_for_pass_through(model="pt-0", request_kwargs={}) def test_initialize_deployment_for_pass_through_keeps_bedrock_iam_deployment(): @@ -6892,9 +6593,7 @@ def test_initialize_deployment_for_pass_through_keeps_bedrock_iam_deployment(): } ] ) - assert [m["model_info"]["id"] for m in router.get_model_list()] == [ - "bedrock-iam-pt" - ] + assert [m["model_info"]["id"] for m in router.get_model_list()] == ["bedrock-iam-pt"] def test_pass_through_deployment_api_key_resolves_via_get_credentials(): @@ -6905,12 +6604,7 @@ def test_pass_through_deployment_api_key_resolves_via_get_credentials(): router = _router_with_two_pass_through_deployments([False, False]) passthrough_router = PassthroughEndpointRouter(llm_router_getter=lambda: router) assert len(router.get_model_list()) == 2 - assert ( - passthrough_router.get_credentials( - custom_llm_provider="openai", region_name=None - ) - == "sk-fake-for-tests" - ) + assert passthrough_router.get_credentials(custom_llm_provider="openai", region_name=None) == "sk-fake-for-tests" def test_get_deployment_credentials_returns_none_for_blocked_deployment(): @@ -6944,16 +6638,9 @@ def test_is_deployment_blocked_static_helper_reflects_blocked_flag(): # No model_info on deployment object → treated as not blocked assert litellm.Router._is_deployment_blocked(object()) is False missing_blocked = types.SimpleNamespace() + assert litellm.Router._is_deployment_blocked(types.SimpleNamespace(model_info=missing_blocked)) is False assert ( - litellm.Router._is_deployment_blocked( - types.SimpleNamespace(model_info=missing_blocked) - ) - is False - ) - assert ( - litellm.Router._is_deployment_blocked( - types.SimpleNamespace(model_info=types.SimpleNamespace(blocked=True)) - ) + litellm.Router._is_deployment_blocked(types.SimpleNamespace(model_info=types.SimpleNamespace(blocked=True))) is True ) @@ -6993,9 +6680,7 @@ class TestRouterRequestTimeoutPropagation: litellm.request_timeout = original_value litellm.request_timeout_explicitly_set = original_flag - def test_request_timeout_stored_independently_when_both_set( - self, explicit_request_timeout - ): + def test_request_timeout_stored_independently_when_both_set(self, explicit_request_timeout): router = self._make_router(timeout=330) assert router.timeout == 330 assert router.request_timeout == 300 @@ -7013,22 +6698,16 @@ class TestRouterRequestTimeoutPropagation: litellm.request_timeout = original_value litellm.request_timeout_explicitly_set = original_flag - def test_non_stream_prefers_request_timeout_over_router_timeout( - self, explicit_request_timeout - ): + def test_non_stream_prefers_request_timeout_over_router_timeout(self, explicit_request_timeout): router = self._make_router(timeout=330) assert router._get_non_stream_timeout(kwargs={}, data={}) == 300 - def test_stream_prefers_request_timeout_over_router_timeout( - self, explicit_request_timeout - ): + def test_stream_prefers_request_timeout_over_router_timeout(self, explicit_request_timeout): router = self._make_router(timeout=330) # stream=True resolves through _get_stream_timeout; request_timeout must win. assert router._get_timeout(kwargs={"stream": True}, data={}) == 300 - def test_explicit_stream_timeout_still_wins_over_request_timeout( - self, explicit_request_timeout - ): + def test_explicit_stream_timeout_still_wins_over_request_timeout(self, explicit_request_timeout): router = self._make_router(timeout=330, stream_timeout=45) assert router._get_stream_timeout(kwargs={}, data={}) == 45 @@ -7044,22 +6723,13 @@ class TestRouterRequestTimeoutPropagation: litellm.request_timeout = original_value litellm.request_timeout_explicitly_set = original_flag - def test_per_deployment_timeout_overrides_request_timeout( - self, explicit_request_timeout - ): + def test_per_deployment_timeout_overrides_request_timeout(self, explicit_request_timeout): router = self._make_router(timeout=330) assert router._get_non_stream_timeout(kwargs={}, data={"timeout": 120}) == 120 - def test_per_request_timeout_overrides_request_timeout( - self, explicit_request_timeout - ): + def test_per_request_timeout_overrides_request_timeout(self, explicit_request_timeout): router = self._make_router(timeout=330) - assert ( - router._get_non_stream_timeout( - kwargs={"timeout": 60}, data={"timeout": 120} - ) - == 60 - ) + assert router._get_non_stream_timeout(kwargs={"timeout": 60}, data={"timeout": 120}) == 60 # --------------------------------------------------------------------------- @@ -7368,9 +7038,7 @@ class TestAdvisorSubCallCooldown: ) def _cooled_down_ids(self, router): - active = router.cooldown_cache.get_active_cooldowns( - model_ids=["dep-1"], parent_otel_span=None - ) + active = router.cooldown_cache.get_active_cooldowns(model_ids=["dep-1"], parent_otel_span=None) return [entry[0] for entry in active] @pytest.mark.asyncio @@ -7379,12 +7047,7 @@ class TestAdvisorSubCallCooldown: router = self._router() now = datetime.now() - assert ( - router.deployment_callback_on_failure( - self._kwargs(self._auth_error()), None, now, now - ) - is True - ) + assert router.deployment_callback_on_failure(self._kwargs(self._auth_error()), None, now, now) is True assert "dep-1" in self._cooled_down_ids(router) def test_advisor_orchestration_failure_does_not_cool_down_deployment(self): @@ -7399,12 +7062,7 @@ class TestAdvisorSubCallCooldown: mark_advisor_orchestration_failure(exception) now = datetime.now() - assert ( - router.deployment_callback_on_failure( - self._kwargs(exception), None, now, now - ) - is False - ) + assert router.deployment_callback_on_failure(self._kwargs(exception), None, now, now) is False assert "dep-1" not in self._cooled_down_ids(router) @@ -7461,13 +7119,13 @@ def test_stream_chunks_have_generated_content_detects_text_and_non_text(): audio_chunk = _chunk(audio_delta) assert _stream_chunks_have_generated_content([audio_chunk]) is True - images_delta = Delta(images=[{"image_url": {"url": "https://example.com/img.png"}, "index": 0, "type": "image_url"}]) + images_delta = Delta( + images=[{"image_url": {"url": "https://example.com/img.png"}, "index": 0, "type": "image_url"}] + ) images_chunk = _chunk(images_delta) assert _stream_chunks_have_generated_content([images_chunk]) is True - annotations_delta = Delta( - annotations=[{"type": "url_citation", "url_citation": {"url": "https://example.com"}}] - ) + annotations_delta = Delta(annotations=[{"type": "url_citation", "url_citation": {"url": "https://example.com"}}]) annotations_chunk = _chunk(annotations_delta) assert _stream_chunks_have_generated_content([annotations_chunk]) is True @@ -7511,12 +7169,8 @@ def test_get_configured_token_limits_skips_wildcard_pattern_matching(): ] ) - with patch.object( - router.pattern_router, "route", side_effect=AssertionError("pattern route called") - ): - assert router.get_configured_token_limits( - "bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0" - ) == (None, None) + with patch.object(router.pattern_router, "route", side_effect=AssertionError("pattern route called")): + assert router.get_configured_token_limits("bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0") == (None, None) def test_get_configured_token_limits_treats_malformed_values_as_absent(): @@ -7589,13 +7243,8 @@ def test_get_configured_mode_skips_wildcard_pattern_matching(): ] ) - with patch.object( - router.pattern_router, "route", side_effect=AssertionError("pattern route called") - ): - assert ( - router.get_configured_mode("bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0") - is None - ) + with patch.object(router.pattern_router, "route", side_effect=AssertionError("pattern route called")): + assert router.get_configured_mode("bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0") is None def test_get_configured_mode_treats_malformed_values_as_absent(): @@ -7654,13 +7303,8 @@ def test_get_configured_display_name_skips_wildcard_pattern_matching(): ] ) - with patch.object( - router.pattern_router, "route", side_effect=AssertionError("pattern route called") - ): - assert ( - router.get_configured_display_name("bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0") - is None - ) + with patch.object(router.pattern_router, "route", side_effect=AssertionError("pattern route called")): + assert router.get_configured_display_name("bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0") is None def test_get_configured_display_name_treats_malformed_values_as_absent(): @@ -7893,13 +7537,16 @@ async def test_acreate_batch_request_bedrock_tags_override_deployment_tags(): mock_client = MagicMock() mock_client.post = AsyncMock(side_effect=lambda *args, **kwargs: fake_response()) - with patch.object( - CommonBatchFilesUtils, - "sign_aws_request", - return_value=({"Authorization": "signed"}, b"{}"), - ) as mock_sign, patch( - "litellm.llms.custom_httpx.llm_http_handler.get_async_httpx_client", - return_value=mock_client, + with ( + patch.object( + CommonBatchFilesUtils, + "sign_aws_request", + return_value=({"Authorization": "signed"}, b"{}"), + ) as mock_sign, + patch( + "litellm.llms.custom_httpx.llm_http_handler.get_async_httpx_client", + return_value=mock_client, + ), ): await router.acreate_batch( model="bedrock-batch-model", @@ -7928,13 +7575,13 @@ async def test_avector_store_search_injects_router(): """ from litellm.types.vector_stores import VectorStoreSearchResponse - expected_response = VectorStoreSearchResponse( - object="vector_store.search_results.page", search_query="q", data=[] - ) + expected_response = VectorStoreSearchResponse(object="vector_store.search_results.page", search_query="q", data=[]) mock_asearch = AsyncMock(return_value=expected_response) # Router.__init__ binds asearch via a local import, so patch the module # attribute before constructing the Router. - with patch("litellm.vector_stores.main.asearch", new=mock_asearch): # test-quality-ok: the SDK call is the only place the injected router kwarg is observable + with patch( + "litellm.vector_stores.main.asearch", new=mock_asearch + ): # test-quality-ok: the SDK call is the only place the injected router kwarg is observable router = litellm.Router( model_list=[ { @@ -7968,7 +7615,9 @@ async def test_avector_store_create_does_not_inject_router(): } ] ) - with patch("litellm.vector_stores.main.acreate", new=mock_acreate): # test-quality-ok: the SDK call is the only place a leaked router kwarg would surface + with patch( + "litellm.vector_stores.main.acreate", new=mock_acreate + ): # test-quality-ok: the SDK call is the only place a leaked router kwarg would surface create_response = await router.avector_store_create(model=None, custom_llm_provider="openai") assert create_response is expected_response @@ -7984,13 +7633,13 @@ def test_vector_store_search_injects_router(): """ from litellm.types.vector_stores import VectorStoreSearchResponse - expected_response = VectorStoreSearchResponse( - object="vector_store.search_results.page", search_query="q", data=[] - ) + expected_response = VectorStoreSearchResponse(object="vector_store.search_results.page", search_query="q", data=[]) mock_search = MagicMock(return_value=expected_response) # Router.__init__ binds search via a local import, so patch the module # attribute before constructing the Router. - with patch("litellm.vector_stores.main.search", new=mock_search): # test-quality-ok: the SDK call is the only place the injected router kwarg is observable + with patch( + "litellm.vector_stores.main.search", new=mock_search + ): # test-quality-ok: the SDK call is the only place the injected router kwarg is observable router = litellm.Router( model_list=[ { @@ -7999,9 +7648,7 @@ def test_vector_store_search_injects_router(): } ] ) - search_response = router.vector_store_search( - vector_store_id="v", query="q", custom_llm_provider="s3_vectors" - ) + search_response = router.vector_store_search(vector_store_id="v", query="q", custom_llm_provider="s3_vectors") assert search_response is expected_response mock_search.assert_called_once() @@ -8015,7 +7662,9 @@ def test_vector_store_create_does_not_inject_router(): mock_create = MagicMock(return_value=expected_response) # Router.__init__ binds create via a local import, so patch the module # attribute before constructing the Router. - with patch("litellm.vector_stores.main.create", new=mock_create): # test-quality-ok: the SDK call is the only place a leaked router kwarg would surface + with patch( + "litellm.vector_stores.main.create", new=mock_create + ): # test-quality-ok: the SDK call is the only place a leaked router kwarg would surface router = litellm.Router( model_list=[ { @@ -8050,9 +7699,7 @@ class TestPreRoutingStrategyRegistryLifecycle: def _complexity_router_params(default_model: str, tags=None) -> dict: return { "model": "auto_router/complexity_router", - "complexity_router_config": { - "tiers": {"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o", "COMPLEX": "gpt-4o"} - }, + "complexity_router_config": {"tiers": {"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o", "COMPLEX": "gpt-4o"}}, "complexity_router_default_model": default_model, **({"tags": tags} if tags else {}), } @@ -8357,9 +8004,7 @@ class TestPreRoutingStrategyRegistryLifecycle: deployment=Deployment( model_name="hybrid-router", litellm_params=LiteLLM_Params( - **self._hybrid_router_params( - {"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o", "COMPLEX": "gpt-4o"} - ) + **self._hybrid_router_params({"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o", "COMPLEX": "gpt-4o"}) ), model_info=ModelInfo(id="router-1", db_model=True), ) @@ -8468,9 +8113,7 @@ class TestPreRoutingStrategyRegistryLifecycle: ({"model": "openai/gpt-4o"}, False), ] for params, expected in cases: - actual = router._deployment_participates_in_adaptive_routing( - litellm_params=LiteLLM_Params(**params) - ) + actual = router._deployment_participates_in_adaptive_routing(litellm_params=LiteLLM_Params(**params)) assert actual is expected, params["model"] @@ -8657,22 +8300,16 @@ class TestUpsertDeploymentRollback: router.delete_deployment(id="prod-1") assert router.has_model_id("prod-1") is False - router._restore_deployment_after_failed_upsert( - previous_deployment=previous, model_id="prod-1" - ) + router._restore_deployment_after_failed_upsert(previous_deployment=previous, model_id="prod-1") restored = router.get_deployment(model_id="prod-1") assert restored is not None assert restored.litellm_params.model == "gpt-4o" - router._restore_deployment_after_failed_upsert( - previous_deployment=previous, model_id="prod-1" - ) + router._restore_deployment_after_failed_upsert(previous_deployment=previous, model_id="prod-1") assert len(router.model_list) == 1 - router._restore_deployment_after_failed_upsert( - previous_deployment=None, model_id="prod-1" - ) + router._restore_deployment_after_failed_upsert(previous_deployment=None, model_id="prod-1") assert len(router.model_list) == 1 @@ -9438,18 +9075,14 @@ class TestAutoRouterSharedModelNameConnectionParams: return httpx.Response( status_code=200, json={ - "candidates": [ - {"content": {"parts": [{"text": "Paris"}], "role": "model"}, "finishReason": "STOP"} - ], + "candidates": [{"content": {"parts": [{"text": "Paris"}], "role": "model"}, "finishReason": "STOP"}], "usageMetadata": {"promptTokenCount": 5, "candidatesTokenCount": 1, "totalTokenCount": 6}, "modelVersion": "gemini-3.6-flash", }, request=httpx.Request("POST", "https://generativelanguage.googleapis.com"), ) - @pytest.mark.parametrize( - "plain_entry_first", [True, False], ids=["plain_entry_first", "marker_entry_first"] - ) + @pytest.mark.parametrize("plain_entry_first", [True, False], ids=["plain_entry_first", "marker_entry_first"]) async def test_routed_tier_call_goes_out_on_its_own_endpoint_and_credentials(self, plain_entry_first): """The outbound provider request for the routed tier hits the tier's own Gemini host with the tier's own key, never the plain sibling's api_base or api_key.""" @@ -9572,9 +9205,7 @@ async def _drive_cyclic_fallback(router, capture, recorder=None, **request_kwarg litellm.callbacks.append(recorder) try: with pytest.raises(litellm.InternalServerError): - await router.acompletion( - model="group-a", messages=[{"role": "user", "content": "hi"}], **request_kwargs - ) + await router.acompletion(model="group-a", messages=[{"role": "user", "content": "hi"}], **request_kwargs) finally: router_logger.removeHandler(capture) router_logger.setLevel(previous_level) @@ -9613,9 +9244,9 @@ async def test_retry_breadcrumbs_do_not_carry_the_walk_state(): breadcrumbs = [breadcrumb for hop in recorder.breadcrumbs_per_target for breadcrumb in hop] assert breadcrumbs, "no retry breadcrumbs were recorded" - assert any( - "fallback_depth" in breadcrumb for breadcrumb in breadcrumbs - ), "no breadcrumb carried router walk state, so this test cannot see the leak" + assert any("fallback_depth" in breadcrumb for breadcrumb in breadcrumbs), ( + "no breadcrumb carried router walk state, so this test cannot see the leak" + ) for breadcrumb in breadcrumbs: assert "attempted_targets" not in breadcrumb @@ -9661,7 +9292,9 @@ async def test_retry_breadcrumbs_never_carry_a_forwarded_credential(container_ke breadcrumbs = metadata["previous_models"] assert breadcrumbs, "no retry breadcrumbs were recorded" dumped = json.dumps(breadcrumbs, default=str) - assert container_key in dumped, "the credential-bearing kwarg never reached the breadcrumb, so this test cannot see the leak" + assert container_key in dumped, ( + "the credential-bearing kwarg never reached the breadcrumb, so this test cannot see the leak" + ) assert _BREADCRUMB_CREDENTIAL_CANARY not in dumped @@ -9766,9 +9399,7 @@ async def test_fallback_failure_detail_from_upstream_is_bounded(): await _drive_cyclic_fallback( _cyclic_fallback_router(), capture, - mock_response=litellm.InternalServerError( - message=huge_message, llm_provider="openai", model="group-a" - ), + mock_response=litellm.InternalServerError(message=huge_message, llm_provider="openai", model="group-a"), ) assert capture.messages, "the fallback failure path did not log at ERROR" @@ -9834,9 +9465,7 @@ def test_ensure_deployment_affinity_callback_is_idempotent(): try: router._ensure_deployment_affinity_callback() router._ensure_deployment_affinity_callback() - affinity_callbacks = [ - cb for cb in router.optional_callbacks or [] if isinstance(cb, DeploymentAffinityCheck) - ] + affinity_callbacks = [cb for cb in router.optional_callbacks or [] if isinstance(cb, DeploymentAffinityCheck)] assert len(affinity_callbacks) == 1 finally: for cb in router.optional_callbacks or []: @@ -9978,9 +9607,7 @@ class TestModelGroupAliasReachesPreRoutingStrategies: router = self._router("auto_routers") metadata: dict = {} - response = await router.acompletion( - model="smart-alias", messages=self._messages(), metadata=metadata - ) + response = await router.acompletion(model="smart-alias", messages=self._messages(), metadata=metadata) assert response.choices[0].message.content == "routed by the tier" assert metadata["model_group"] == "smart-alias" @@ -10024,9 +9651,7 @@ class TestAutoRouterCompressionDecoupling: 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 - ] + compressed = [{**m, "content": f"[COMPRESSED] {m.get('content')}"} for m in structured_messages] return {**inputs, "structured_messages": compressed} @staticmethod @@ -10105,9 +9730,7 @@ class TestAutoRouterCompressionDecoupling: 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 - ): + 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. @@ -10131,11 +9754,37 @@ class TestAutoRouterCompressionDecoupling: 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", @@ -10145,13 +9794,17 @@ class TestAutoRouterCompressionDecoupling: # 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?"} - ] + 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) - response = await router.async_pre_routing_hook( - model="smart-router", request_kwargs={"metadata": {}}, messages=already_compressed_messages - ) + 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 @@ -10197,17 +9850,14 @@ class TestAzureBaseModelFallbackLogging: def test_map_known_deployment_name_resolves_without_error_log(self): router = self._router_with_azure_deployment("azure/gpt-4o") - with patch( - "litellm.router.verbose_router_logger.error" - ) as mock_error: + with patch("litellm.router.verbose_router_logger.error") as mock_error: model_info = router.get_router_model_info( deployment=None, received_model_name="my-group", id="azure-base-model-test-id" ) - assert not any( - "Could not identify azure model" in str(call) - for call in mock_error.call_args_list - ), f"unexpected error log: {mock_error.call_args_list}" + assert not any("Could not identify azure model" in str(call) for call in mock_error.call_args_list), ( + f"unexpected error log: {mock_error.call_args_list}" + ) # the fallback resolution must actually surface the map values assert model_info["max_input_tokens"] == litellm.model_cost["azure/gpt-4o"]["max_input_tokens"] assert model_info["input_cost_per_token"] == litellm.model_cost["azure/gpt-4o"]["input_cost_per_token"] @@ -10215,17 +9865,14 @@ class TestAzureBaseModelFallbackLogging: def test_unmappable_deployment_name_still_logs_error(self): router = self._router_with_azure_deployment("azure/my-custom-deployment-name") - with patch( - "litellm.router.verbose_router_logger.error" - ) as mock_error: + with patch("litellm.router.verbose_router_logger.error") as mock_error: model_info = router.get_router_model_info( deployment=None, received_model_name="my-group", id="azure-base-model-test-id" ) - assert any( - "Could not identify azure model" in str(call) - for call in mock_error.call_args_list - ), "expected the error log for an unmappable azure deployment name" + assert any("Could not identify azure model" in str(call) for call in mock_error.call_args_list), ( + "expected the error log for an unmappable azure deployment name" + ) # unmappable names resolve to a zeroed stub — unchanged behavior assert model_info.get("max_input_tokens") is None @@ -10252,6 +9899,7 @@ class TestAzureBaseModelFallbackLogging: ) assert model_info["max_input_tokens"] == litellm.model_cost["azure/gpt-4o-mini"]["max_input_tokens"] + def test_model_group_info_intersects_supported_reasoning_efforts(): router = litellm.Router( model_list=[ @@ -10342,7 +9990,6 @@ def test_model_group_info_reasoning_efforts_are_unknown_when_any_deployment_is_o assert result.supported_reasoning_efforts is None - def test_model_group_info_surfaces_supports_parallel_function_calling(local_model_cost_map): """``/model_group/info`` folds each deployment's registry flags into the group; a deployment whose registry entry declares parallel function calling must flip the group to True instead of False.""" @@ -10575,6 +10222,7 @@ class TestAddDeploymentApiBaseProviderResolution: assert deployment is not None assert deployment.litellm_params.custom_llm_provider == "openai" + # ===================================================================== # anthropic_messages mid-stream-fallback helpers, added for #24004 # (mid-stream fallback not supported for anthropic_messages route type). @@ -10685,10 +10333,7 @@ class _AnthropicMessagesFallbackByteStream: def _anthropic_messages_overloaded_error_chunk() -> bytes: - return ( - b"event: error\n" - b'data: {"type": "error", "error": {"type": "overloaded_error", "message": "Overloaded"}}\n\n' - ) + return b'event: error\ndata: {"type": "error", "error": {"type": "overloaded_error", "message": "Overloaded"}}\n\n' def _anthropic_messages_invalid_request_error_chunk() -> bytes: @@ -10733,9 +10378,7 @@ async def test_anthropic_messages_streaming_iterator_passthrough(): [_anthropic_messages_content_chunk("hi"), _anthropic_messages_content_chunk(" there")] ) - wrapped = await router._aanthropic_messages_streaming_iterator( - response=source, initial_kwargs={"model": "primary"} - ) + wrapped = await router._aanthropic_messages_streaming_iterator(response=source, initial_kwargs={"model": "primary"}) collected = [chunk async for chunk in wrapped] assert collected == [_anthropic_messages_content_chunk("hi"), _anthropic_messages_content_chunk(" there")] @@ -10754,12 +10397,14 @@ async def test_anthropic_messages_streaming_iterator_flushes_buffered_lifecycle_ [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("hi"), message_stop] ) - wrapped = await router._aanthropic_messages_streaming_iterator( - response=source, initial_kwargs={"model": "primary"} - ) + wrapped = await router._aanthropic_messages_streaming_iterator(response=source, initial_kwargs={"model": "primary"}) collected = [chunk async for chunk in wrapped] - assert collected == [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("hi"), message_stop] + assert collected == [ + _anthropic_messages_message_start_chunk(), + _anthropic_messages_content_chunk("hi"), + message_stop, + ] @pytest.mark.asyncio @@ -10771,9 +10416,7 @@ async def test_anthropic_messages_streaming_iterator_flushes_buffered_frames_on_ message_stop = b'event: message_stop\ndata: {"type": "message_stop"}\n\n' source = _AnthropicMessagesFakeByteStream([_anthropic_messages_message_start_chunk(), message_stop]) - wrapped = await router._aanthropic_messages_streaming_iterator( - response=source, initial_kwargs={"model": "primary"} - ) + wrapped = await router._aanthropic_messages_streaming_iterator(response=source, initial_kwargs={"model": "primary"}) collected = [chunk async for chunk in wrapped] assert collected == [_anthropic_messages_message_start_chunk(), message_stop] @@ -10844,7 +10487,9 @@ async def test_anthropic_messages_leading_ping_keepalive_is_forwarded_live(): yield _anthropic_messages_message_start_chunk() yield _anthropic_messages_content_chunk("hi") - wrapped = await router._aanthropic_messages_streaming_iterator(response=source(), initial_kwargs={"model": "primary"}) + wrapped = await router._aanthropic_messages_streaming_iterator( + response=source(), initial_kwargs={"model": "primary"} + ) assert await asyncio.wait_for(wrapped.__anext__(), timeout=1) == _anthropic_messages_ping_chunk() content_released.set() @@ -12645,9 +12290,9 @@ class TestTierParamsTheTargetAccepts: def test_declared_param_allowlist_ignores_malformed_declarations(self): """A str is iterable, so without the type guard a YAML scalar mistake like allowed_openai_params: reasoning_effort would allowlist single characters.""" - assert litellm.Router._declared_param_allowlist({"allowed_openai_params": ["reasoning_effort", 3]}) == frozenset( - {"reasoning_effort"} - ) + assert litellm.Router._declared_param_allowlist( + {"allowed_openai_params": ["reasoning_effort", 3]} + ) == frozenset({"reasoning_effort"}) assert litellm.Router._declared_param_allowlist({"allowed_openai_params": "reasoning_effort"}) == frozenset() assert litellm.Router._declared_param_allowlist({}) == frozenset() @@ -12727,7 +12372,11 @@ class TestTierParamsTheTargetAccepts: @pytest.mark.parametrize( "deployment", - [{"model_name": "x"}, {"model_name": "x", "litellm_params": {}}, {"model_name": "x", "litellm_params": {"model": "not-a-real-provider/nope"}}], + [ + {"model_name": "x"}, + {"model_name": "x", "litellm_params": {}}, + {"model_name": "x", "litellm_params": {"model": "not-a-real-provider/nope"}}, + ], ) def test_deployment_accepts_param_fails_open(self, deployment): """An unresolvable deployment must not be the reason a param is dropped.""" @@ -12956,9 +12605,7 @@ class TestPreRoutingTierDrivesFallbacks: async def test_the_selected_tier_fallback_chain_runs(self): router = self._router([{"tier1": ["backup-a"]}]) - response = await router.acompletion( - model="smart-router", messages=[{"role": "user", "content": "hi"}] - ) + response = await router.acompletion(model="smart-router", messages=[{"role": "user", "content": "hi"}]) assert response.choices[0].message.content == "from backup-a" @@ -12991,9 +12638,7 @@ class TestPreRoutingTierDrivesFallbacks: async def test_a_request_without_a_pre_routing_hook_still_uses_its_own_group(self): router = self._router([{"tier1": ["backup-a"]}]) - response = await router.acompletion( - model="tier1", messages=[{"role": "user", "content": "hi"}] - ) + response = await router.acompletion(model="tier1", messages=[{"role": "user", "content": "hi"}]) assert response.choices[0].message.content == "from backup-a" diff --git a/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.test.ts b/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.test.ts index b917fcedaa2..ea6eaf99106 100644 --- a/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.test.ts +++ b/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.test.ts @@ -80,9 +80,20 @@ describe("hydrateAutoRouterCompression", () => { expect(state).toEqual({ routing: "headroom-a", sameAsRouting: false, model: "headroom-b" }); }); - it("treats a missing model key as same-as-routing", () => { + 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: true, model: undefined }); + 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("round-trips through buildAutoRouterCompressionParams", () => { diff --git a/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.ts b/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.ts index c86416b507f..47d0e3db7e4 100644 --- a/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.ts +++ b/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.ts @@ -53,7 +53,11 @@ export const hydrateAutoRouterCompression = (litellmParams: { const routing = litellmParams.auto_router_routing_compression ?? undefined; if (routing === undefined) return DEFAULT_AUTO_ROUTER_COMPRESSION; - const model = litellmParams.auto_router_model_compression ?? undefined; - const sameAsRouting = model === undefined || model === routing; + // An absent model key is no model-hop compression, not same-as-routing: the backend + // reads it as None (policy_from_litellm_params). Hydrating it as same-as-routing + // would make re-saving an unrelated edit write the routing guardrail onto the model + // hop and silently start compressing the model call. + const model = litellmParams.auto_router_model_compression ?? NO_COMPRESSION; + const sameAsRouting = model === routing; return { routing, sameAsRouting, model: sameAsRouting ? undefined : model }; }; From 8284208af261bd32d78ff3fb43040117894fd358 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Sat, 5 Sep 2026 09:51:51 -0700 Subject: [PATCH 09/13] fix(auto-router compression): restrict both hops to real compression guardrails The two policy fields are operator-supplied names and nothing else constrained them. The routing hop calls apply_guardrail directly, which hands the guardrail the conversation and POSTs it to whatever service backs that guardrail, and the model hop is added to metadata["guardrails"], which runs it even when it is not default_on. So naming an ordinary guardrail turned either hop into a way to invoke it and ship prompt content to it. Both hops now refuse a name that does not resolve to an active compression guardrail, and say so in the log rather than failing quietly. --- .../guardrails/auto_router_compression.py | 52 ++++++++++++++--- .../test_auto_router_compression.py | 56 +++++++++++++++++-- tests/test_litellm/test_router.py | 7 ++- 3 files changed, 103 insertions(+), 12 deletions(-) diff --git a/litellm/proxy/guardrails/auto_router_compression.py b/litellm/proxy/guardrails/auto_router_compression.py index e7f58662249..b6f32c1c46d 100644 --- a/litellm/proxy/guardrails/auto_router_compression.py +++ b/litellm/proxy/guardrails/auto_router_compression.py @@ -135,19 +135,36 @@ def team_id_from_request(request_kwargs: Mapping[str, object]) -> str | None: 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 are validated through here. The two policy fields are operator-supplied + names and nothing else constrains them, so without this a name that resolves to an + ordinary guardrail would be handed the conversation and invoked: the routing hop + calls `apply_guardrail` directly, which POSTs the content wherever that guardrail + sends it, and the model hop is added to `metadata["guardrails"]`, which runs it even + when it is not `default_on`. + """ + 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 - from litellm.proxy.guardrails.guardrail_registry import guardrail_class_registry - compression_classes: Final = tuple( - cls for name, cls in guardrail_class_registry.items() if name in COMPRESSION_GUARDRAIL_PROVIDERS - ) - if not compression_classes: + 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 isinstance(cb, compression_classes) and cb.guardrail_name) + return tuple(cb for cb in active if is_compression_guardrail(cb) and cb.guardrail_name) async def arm_pre_call( @@ -192,7 +209,18 @@ async def arm_pre_call( ) ) - if policy.model is not None: + # Only a name that resolves to a real compression guardrail may be armed: this adds + # it to `metadata["guardrails"]`, which runs it even when it is 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") @@ -249,6 +277,16 @@ async def messages_for_routing( ) return _as_routing_messages(messages) + # apply_guardrail below hands this guardrail the conversation and it POSTs the + # content to whatever service backs it, so the name has to be a compression + # guardrail rather than any guardrail the operator happened to name. + 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 } diff --git a/tests/test_litellm/proxy/guardrails/test_auto_router_compression.py b/tests/test_litellm/proxy/guardrails/test_auto_router_compression.py index 0b47e56cb02..676b3ba2967 100644 --- a/tests/test_litellm/proxy/guardrails/test_auto_router_compression.py +++ b/tests/test_litellm/proxy/guardrails/test_auto_router_compression.py @@ -186,15 +186,33 @@ class _RecordingCompressionGuardrail(CustomGuardrail): @pytest.fixture -def registered_guardrail(): +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): @@ -266,7 +284,13 @@ class TestArmPreCall: 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): + 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( [ { @@ -280,8 +304,11 @@ class TestArmPreCall: ] ) data = {"model": "smart-router", "messages": [{"role": "user", "content": "hi"}]} - await arm_pre_call(data=data, llm_router=router) - assert data["metadata"]["guardrails"] == ["headroom-b"] + 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): @@ -347,6 +374,27 @@ class TestMessagesForRouting: 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): the policy fields are operator-supplied names that + nothing else constrains. apply_guardrail hands the guardrail the conversation + and it POSTs that content to whatever service backs it, so naming an ordinary + guardrail must not turn the routing hop into a way to ship prompts there.""" + 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 real compression guardrail writes its stats onto whatever diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 9e6a88d3433..6279c8a5404 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -9686,7 +9686,12 @@ class TestAutoRouterCompressionDecoupling: return router, strategy @pytest.fixture - def registered_guardrail(self): + 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 From 1c16a5910b07415b2ab9cb6a54622f0296117f93 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Sat, 5 Sep 2026 10:31:56 -0700 Subject: [PATCH 10/13] fix(tests): undo a stray whole-file reformat and arm a real guardrail test_router.py is not ruff-formatted on staging and CI's format check only scopes litellm/*.py, so running ruff format over the whole file rewrote ~900 lines of unrelated code. That reflow split long single-line patch() calls into multi-line form, which the test-quality gate counts individually, pushing TQ008 four over its ceiling. The file is back to staging's formatting with only the compression test class added. test_common_request_processing.py armed a model-side guardrail name with no such guardrail registered, which stopped working once both hops began requiring the name to resolve to an active compression guardrail. --- .../proxy/test_common_request_processing.py | 33 +- tests/test_litellm/test_router.py | 919 +++++++++++++----- 2 files changed, 675 insertions(+), 277 deletions(-) diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index c0809e53d2e..96d9b0c5a26 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -383,6 +383,18 @@ class TestProxyBaseLLMRequestProcessing: """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 = {} @@ -418,15 +430,18 @@ class TestProxyBaseLLMRequestProcessing: mock_proxy_config = MagicMock(spec=ProxyConfig) mock_proxy_config._get_hierarchical_router_settings = AsyncMock(return_value=None) - await processing_obj.common_processing_pre_call_logic( - request=mock_request, - general_settings={}, - user_api_key_dict=ProxyUserAPIKeyAuth(api_key="sk-test"), - proxy_logging_obj=mock_proxy_logging_obj, - proxy_config=mock_proxy_config, - route_type="acompletion", - llm_router=fake_llm_router, - ) + 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"] diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 6279c8a5404..d97fa7f2912 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -15,6 +15,7 @@ import pytest import respx + import litellm from litellm import Router from litellm.exceptions import MidStreamFallbackError @@ -136,18 +137,31 @@ def test_router_model_group_encrypted_content_affinity_callback_registration(): num_retries=0, ) callbacks = router.optional_callbacks or [] - encrypted_content_callbacks = [cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck)] - deployment_callback = next(cb for cb in callbacks if isinstance(cb, DeploymentAffinityCheck)) + encrypted_content_callbacks = [ + cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck) + ] + deployment_callback = next( + cb for cb in callbacks if isinstance(cb, DeploymentAffinityCheck) + ) assert len(encrypted_content_callbacks) == 1 assert encrypted_content_callbacks[0].enable_global_affinity is False - assert encrypted_content_callbacks[0].model_group_affinity_config == model_group_affinity_config - assert callbacks.index(encrypted_content_callbacks[0]) < callbacks.index(deployment_callback) - assert litellm.callbacks.index(encrypted_content_callbacks[0]) < (litellm.callbacks.index(deployment_callback)) + assert ( + encrypted_content_callbacks[0].model_group_affinity_config + == model_group_affinity_config + ) + assert callbacks.index(encrypted_content_callbacks[0]) < callbacks.index( + deployment_callback + ) + assert litellm.callbacks.index(encrypted_content_callbacks[0]) < ( + litellm.callbacks.index(deployment_callback) + ) router._add_encrypted_content_affinity_check(enable_global_affinity=True) callbacks = router.optional_callbacks or [] - encrypted_content_callbacks = [cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck)] + encrypted_content_callbacks = [ + cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck) + ] assert len(encrypted_content_callbacks) == 1 assert encrypted_content_callbacks[0].enable_global_affinity is True assert encrypted_content_callbacks[0].router is router @@ -178,9 +192,13 @@ async def test_encrypted_content_affinity_model_group_config_is_additive(): }, target_deployment, ] - encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id("deployment-b", "rs_test") + encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id( + "deployment-b", "rs_test" + ) - assert EncryptedContentAffinityCheck.has_model_group_affinity_enabled({model_group: ["encrypted_content_affinity"]}) + assert EncryptedContentAffinityCheck.has_model_group_affinity_enabled( + {model_group: ["encrypted_content_affinity"]} + ) assert not EncryptedContentAffinityCheck.has_model_group_affinity_enabled(None) per_group_check = EncryptedContentAffinityCheck( @@ -221,7 +239,10 @@ async def test_encrypted_content_affinity_model_group_config_is_additive(): ) assert unfiltered == healthy_deployments - assert "encrypted_content_affinity_enabled" not in disabled_request_kwargs["litellm_metadata"] + assert ( + "encrypted_content_affinity_enabled" + not in disabled_request_kwargs["litellm_metadata"] + ) global_check = EncryptedContentAffinityCheck( enable_global_affinity=True, @@ -241,7 +262,9 @@ async def test_encrypted_content_affinity_model_group_config_is_additive(): ) assert globally_filtered == [target_deployment] - assert global_request_kwargs["litellm_metadata"]["encrypted_content_affinity_enabled"] + assert global_request_kwargs["litellm_metadata"][ + "encrypted_content_affinity_enabled" + ] @pytest.mark.asyncio @@ -288,10 +311,18 @@ async def test_encrypted_content_affinity_takes_priority_over_user_key_affinity( num_retries=0, ) callbacks = router.optional_callbacks or [] - deployment_callback = next(cb for cb in callbacks if isinstance(cb, DeploymentAffinityCheck)) - encrypted_content_callback = next(cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck)) - assert callbacks.index(encrypted_content_callback) < callbacks.index(deployment_callback) - assert litellm.callbacks.index(encrypted_content_callback) < (litellm.callbacks.index(deployment_callback)) + deployment_callback = next( + cb for cb in callbacks if isinstance(cb, DeploymentAffinityCheck) + ) + encrypted_content_callback = next( + cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck) + ) + assert callbacks.index(encrypted_content_callback) < callbacks.index( + deployment_callback + ) + assert litellm.callbacks.index(encrypted_content_callback) < ( + litellm.callbacks.index(deployment_callback) + ) cache_key = DeploymentAffinityCheck.get_affinity_cache_key( model_group=model_group, @@ -302,7 +333,9 @@ async def test_encrypted_content_affinity_takes_priority_over_user_key_affinity( value={"model_id": "deployment-a"}, ttl=60, ) - encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id("deployment-b", "rs_test") + encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id( + "deployment-b", "rs_test" + ) request_kwargs = { "input": [{"type": "reasoning", "id": encoded_id}], "litellm_metadata": {"user_api_key_hash": user_api_key_hash}, @@ -906,7 +939,9 @@ async def test_arouter_aretrieve_batch(): ], ) - with patch.object(litellm, "aretrieve_batch", return_value=AsyncMock()) as mock_aretrieve_batch: + with patch.object( + litellm, "aretrieve_batch", return_value=AsyncMock() + ) as mock_aretrieve_batch: try: response = await router.aretrieve_batch( model="gpt-3.5-turbo", @@ -927,7 +962,9 @@ async def test_arouter_aretrieve_file_content(): Test that router.acreate_file with JSONL file returns the correct response """ - with patch.object(litellm, "afile_content", return_value=AsyncMock()) as mock_afile_content: + with patch.object( + litellm, "afile_content", return_value=AsyncMock() + ) as mock_afile_content: router = litellm.Router( model_list=[ { @@ -988,7 +1025,7 @@ async def test_arouter_filter_team_based_models(): assert result is not None # FAILS - with pytest.raises(Exception, match="No deployments available for selected model, Try again in") as e: + with pytest.raises(Exception, match='No deployments available for selected model, Try again in') as e: result = await router.acompletion( model="gpt-3.5-turbo", messages=[{"role": "user", "content": "Hello, world!"}], @@ -1072,7 +1109,9 @@ def test_arouter_should_include_deployment(): model=deployment_with_team_and_public_name, team_id="test-team", ) - assert result is True, "Should return True when team_id and team_public_model_name match" + assert ( + result is True + ), "Should return True when team_id and team_public_model_name match" # Test Case 2: Team-specific deployment - team_id matches but model_name doesn't match team_public_model_name result = router.should_include_deployment( @@ -1080,9 +1119,9 @@ def test_arouter_should_include_deployment(): model=deployment_with_team_and_public_name, team_id="test-team", ) - assert result is False, ( - "Should return False when team_id matches but model_name doesn't match team_public_model_name" - ) + assert ( + result is False + ), "Should return False when team_id matches but model_name doesn't match team_public_model_name" # Test Case 3: Team-specific deployment - team_id doesn't match result = router.should_include_deployment( @@ -1098,18 +1137,30 @@ def test_arouter_should_include_deployment(): model=deployment_with_team_no_public_name, team_id="test-team", ) - assert result is True, "Should return True when team deployment has no team_public_model_name to match" + assert ( + result is True + ), "Should return True when team deployment has no team_public_model_name to match" # Test Case 5: Non-team deployment - model_name matches and no team_id - result = router.should_include_deployment(model_name="gpt-4", model=deployment_without_team, team_id=None) - assert result is True, "Should return True when model_name matches and deployment has no team_id" + result = router.should_include_deployment( + model_name="gpt-4", model=deployment_without_team, team_id=None + ) + assert ( + result is True + ), "Should return True when model_name matches and deployment has no team_id" # Test Case 6: Non-team deployment - model_name matches but team_id provided (should still work) - result = router.should_include_deployment(model_name="gpt-4", model=deployment_without_team, team_id="any-team") - assert result is True, "Should return True when model_name matches non-team deployment, regardless of team_id param" + result = router.should_include_deployment( + model_name="gpt-4", model=deployment_without_team, team_id="any-team" + ) + assert ( + result is True + ), "Should return True when model_name matches non-team deployment, regardless of team_id param" # Test Case 7: Non-team deployment - model_name doesn't match - result = router.should_include_deployment(model_name="different-model", model=deployment_without_team, team_id=None) + result = router.should_include_deployment( + model_name="different-model", model=deployment_without_team, team_id=None + ) assert result is False, "Should return False when model_name doesn't match" # Test Case 8: Team deployment accessed without matching team_id @@ -1118,7 +1169,9 @@ def test_arouter_should_include_deployment(): model=deployment_with_team_and_public_name, team_id=None, ) - assert result is True, "Should return True when matching model with exact model_name" + assert ( + result is True + ), "Should return True when matching model with exact model_name" def test_arouter_responses_api_bridge(): @@ -1168,7 +1221,9 @@ def test_arouter_responses_api_bridge(): "status": "completed", "output": [], } - mock_response.text = '{"id": "resp_test", "object": "response", "status": "completed", "output": []}' + mock_response.text = ( + '{"id": "resp_test", "object": "response", "status": "completed", "output": []}' + ) with patch.object(client, "post", return_value=mock_response) as mock_post: try: @@ -1242,7 +1297,7 @@ def test_add_invalid_provider_to_router(): ], ) - with pytest.raises(Exception, match="Unsupported provider - vertex_ai_eu") as e: + with pytest.raises(Exception, match='Unsupported provider - vertex_ai_eu') as e: router.add_deployment( Deployment( model_name="vertex_ai/*", @@ -1294,9 +1349,15 @@ async def test_router_ageneric_api_call_with_fallbacks_helper(): }, } - with patch.object(router, "_update_kwargs_with_deployment") as mock_update_kwargs: - with patch.object(router, "async_routing_strategy_pre_call_checks") as mock_pre_call_checks: - with patch.object(router, "_get_client", return_value=None) as mock_get_client: + with patch.object( + router, "_update_kwargs_with_deployment" + ) as mock_update_kwargs: + with patch.object( + router, "async_routing_strategy_pre_call_checks" + ) as mock_pre_call_checks: + with patch.object( + router, "_get_client", return_value=None + ) as mock_get_client: result = await router._ageneric_api_call_with_fallbacks_helper( model="gpt-3.5-turbo", original_generic_function=mock_generic_function, @@ -1331,7 +1392,7 @@ async def test_router_ageneric_api_call_with_fallbacks_helper(): with patch.object(router, "async_get_available_deployment") as mock_get_deployment: mock_get_deployment.side_effect = Exception("No deployment available") - with pytest.raises(Exception, match="No deployment available") as exc_info: + with pytest.raises(Exception, match='No deployment available') as exc_info: await router._ageneric_api_call_with_fallbacks_helper( model="gpt-3.5-turbo", original_generic_function=mock_generic_function, @@ -1359,9 +1420,15 @@ async def test_router_ageneric_api_call_with_fallbacks_helper(): mock_semaphore = asyncio.Semaphore(1) - with patch.object(router, "_update_kwargs_with_deployment") as mock_update_kwargs: - with patch.object(router, "_get_client", return_value=mock_semaphore) as mock_get_client: - with patch.object(router, "async_routing_strategy_pre_call_checks") as mock_pre_call_checks: + with patch.object( + router, "_update_kwargs_with_deployment" + ) as mock_update_kwargs: + with patch.object( + router, "_get_client", return_value=mock_semaphore + ) as mock_get_client: + with patch.object( + router, "async_routing_strategy_pre_call_checks" + ) as mock_pre_call_checks: result = await router._ageneric_api_call_with_fallbacks_helper( model="gpt-3.5-turbo", original_generic_function=mock_semaphore_function, @@ -1390,10 +1457,16 @@ async def test_router_ageneric_api_call_with_fallbacks_helper(): }, } - with patch.object(router, "_update_kwargs_with_deployment") as mock_update_kwargs: - with patch.object(router, "_get_client", return_value=None) as mock_get_client: - with patch.object(router, "async_routing_strategy_pre_call_checks") as mock_pre_call_checks: - with pytest.raises(Exception, match="Mock failure") as exc_info: + with patch.object( + router, "_update_kwargs_with_deployment" + ) as mock_update_kwargs: + with patch.object( + router, "_get_client", return_value=None + ) as mock_get_client: + with patch.object( + router, "async_routing_strategy_pre_call_checks" + ) as mock_pre_call_checks: + with pytest.raises(Exception, match='Mock failure') as exc_info: await router._ageneric_api_call_with_fallbacks_helper( model="gpt-3.5-turbo", original_generic_function=mock_failing_function, @@ -1461,9 +1534,9 @@ async def test_ageneric_api_call_deployment_model_overrides_alias(): original_generic_function=capture_model, ) - assert captured["model"] == "vertex_ai/gemini-2.5-flash", ( - f"Expected deployment model 'vertex_ai/gemini-2.5-flash', got '{captured['model']}'" - ) + assert ( + captured["model"] == "vertex_ai/gemini-2.5-flash" + ), f"Expected deployment model 'vertex_ai/gemini-2.5-flash', got '{captured['model']}'" @pytest.mark.asyncio @@ -1569,10 +1642,14 @@ def test_router_get_model_access_groups_team_only_models(): ] ) - access_groups = router.get_model_access_groups(model_name="gpt-3.5-turbo", team_id=None) + access_groups = router.get_model_access_groups( + model_name="gpt-3.5-turbo", team_id=None + ) assert len(access_groups) == 0 - access_groups = router.get_model_access_groups(model_name="gpt-3.5-turbo", team_id="team_1") + access_groups = router.get_model_access_groups( + model_name="gpt-3.5-turbo", team_id="team_1" + ) assert list(access_groups.keys()) == ["default-models"] @@ -1667,7 +1744,9 @@ def test_model_group_info_cost_from_db_model_info(): ] ) - with patch.object(router, "get_deployment_model_info", side_effect=Exception("not found")): + with patch.object( + router, "get_deployment_model_info", side_effect=Exception("not found") + ): result = router._cached_get_model_group_info("my-custom-model") assert result is not None assert result.input_cost_per_token == 0.0001 @@ -1695,7 +1774,9 @@ def test_model_group_info_cost_none_when_db_model_info_has_no_cost(): ] ) - with patch.object(router, "get_deployment_model_info", side_effect=Exception("not found")): + with patch.object( + router, "get_deployment_model_info", side_effect=Exception("not found") + ): result = router._cached_get_model_group_info("my-custom-model-no-cost") assert result is not None assert result.input_cost_per_token is None @@ -1765,7 +1846,9 @@ def test_model_group_info_with_stringified_cost_values(): } return None - with patch.object(router, "get_deployment_model_info", side_effect=_model_info_with_str_costs): + with patch.object( + router, "get_deployment_model_info", side_effect=_model_info_with_str_costs + ): result = router._set_model_group_info( model_group="my-custom-model", user_facing_model_group_name="my-custom-model", @@ -1811,7 +1894,9 @@ def test_model_group_info_db_fallback_with_stringified_cost_values(): ] ) - with patch.object(router, "get_deployment_model_info", side_effect=Exception("not found")): + with patch.object( + router, "get_deployment_model_info", side_effect=Exception("not found") + ): result = router._set_model_group_info( model_group="my-custom-model", user_facing_model_group_name="my-custom-model", @@ -2071,7 +2156,6 @@ async def test_acompletion_streaming_iterator(): # Collect streamed chunks — the first chunk succeeds, then the error re-raises collected_chunks = [] - async def _drain(): async for chunk in result: collected_chunks.append(chunk) @@ -2765,7 +2849,11 @@ def _make_responses_iterator( BaseResponsesAPIStreamingIterator, ) - base = LiteLLMCompletionStreamingIterator if bridge else BaseResponsesAPIStreamingIterator + base = ( + LiteLLMCompletionStreamingIterator + if bridge + else BaseResponsesAPIStreamingIterator + ) class _Iter(base): def __init__(self): @@ -2845,7 +2933,9 @@ async def test_aresponses_streaming_iterator_fallback(): BaseResponsesAPIStreamingIterator, ) - router = _make_router_with_fallback("anthropic/claude-sonnet-4-6", "vertex_ai/claude-sonnet-4-6") + router = _make_router_with_fallback( + "anthropic/claude-sonnet-4-6", "vertex_ai/claude-sonnet-4-6" + ) src = _make_responses_iterator( chunks=[MagicMock(type="response.created")], error=MidStreamFallbackError( @@ -2928,9 +3018,9 @@ async def test_aresponses_streaming_iterator_writes_litellm_metadata_on_fallback fbk = mock_fallback_utils.call_args.kwargs["kwargs"] assert "litellm_metadata" in fbk, "wrong metadata_variable_name" assert fbk["litellm_metadata"]["model_group"] == "gpt-4" - assert "model_group" not in fbk.get("metadata", {}), ( - "model_group leaked into 'metadata' instead of 'litellm_metadata'" - ) + assert "model_group" not in fbk.get( + "metadata", {} + ), "model_group leaked into 'metadata' instead of 'litellm_metadata'" @pytest.mark.asyncio @@ -3048,7 +3138,9 @@ async def test_aresponses_streaming_iterator_combines_partial_usage(): fallback_response_object = ResponsesAPIResponse( id="resp_test", created_at=0, model="gpt-4", object="response", output=[] ) - fallback_response_object.usage = ResponseAPIUsage(input_tokens=20, output_tokens=15, total_tokens=35) + fallback_response_object.usage = ResponseAPIUsage( + input_tokens=20, output_tokens=15, total_tokens=35 + ) fallback_event = ResponseCompletedEvent( type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, response=fallback_response_object, @@ -3057,7 +3149,9 @@ async def test_aresponses_streaming_iterator_combines_partial_usage(): with ( patch( "litellm.main.stream_chunk_builder", - return_value=SimpleNamespace(usage=SimpleNamespace(prompt_tokens=10, completion_tokens=4)), + return_value=SimpleNamespace( + usage=SimpleNamespace(prompt_tokens=10, completion_tokens=4) + ), ), patch.object( router, @@ -3358,7 +3452,9 @@ def test_pre_call_checks_skips_token_count_without_max_input_tokens(monkeypatch) monkeypatch.setattr(router, "get_router_model_info", lambda **kwargs: {}) calls = [] - monkeypatch.setattr(litellm, "token_counter", lambda *a, **k: calls.append(1) or 1000) + monkeypatch.setattr( + litellm, "token_counter", lambda *a, **k: calls.append(1) or 1000 + ) deployments = [ {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, @@ -3386,10 +3482,14 @@ def test_pre_call_checks_counts_once_and_filters_on_max_input_tokens(monkeypatch ], enable_pre_call_checks=True, ) - monkeypatch.setattr(router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 5}) + monkeypatch.setattr( + router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 5} + ) calls = [] - monkeypatch.setattr(litellm, "token_counter", lambda *a, **k: calls.append(1) or 1000) + monkeypatch.setattr( + litellm, "token_counter", lambda *a, **k: calls.append(1) or 1000 + ) deployments = [ {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, @@ -3416,10 +3516,14 @@ def test_pre_call_checks_uses_precounted_tokens(monkeypatch): ], enable_pre_call_checks=True, ) - monkeypatch.setattr(router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 5}) + monkeypatch.setattr( + router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 5} + ) calls = [] - monkeypatch.setattr(litellm, "token_counter", lambda *a, **k: calls.append(1) or 1) + monkeypatch.setattr( + litellm, "token_counter", lambda *a, **k: calls.append(1) or 1 + ) deployments = [ {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, @@ -3446,7 +3550,9 @@ async def test_async_get_healthy_deployments_counts_tokens_off_the_event_loop(mo ], enable_pre_call_checks=True, ) - monkeypatch.setattr(router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 1_000_000}) + monkeypatch.setattr( + router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 1_000_000} + ) counting_threads = [] monkeypatch.setattr( @@ -3542,10 +3648,14 @@ def test_pre_call_checks_does_not_recount_inline_after_an_off_loop_failure(monke ], enable_pre_call_checks=True, ) - monkeypatch.setattr(router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 5}) + monkeypatch.setattr( + router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 5} + ) calls = [] - monkeypatch.setattr(litellm, "token_counter", lambda *a, **k: calls.append(1) or 1000) + monkeypatch.setattr( + litellm, "token_counter", lambda *a, **k: calls.append(1) or 1000 + ) deployments = [ {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, @@ -3573,7 +3683,9 @@ async def test_async_get_healthy_deployments_never_recounts_on_the_loop(monkeypa ], enable_pre_call_checks=True, ) - monkeypatch.setattr(router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 5}) + monkeypatch.setattr( + router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 5} + ) counting_threads = [] @@ -3608,7 +3720,9 @@ async def test_acount_pre_call_check_tokens_leaves_the_event_loop_free(monkeypat ], enable_pre_call_checks=True, ) - monkeypatch.setattr(router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 5}) + monkeypatch.setattr( + router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 5} + ) deployments = [ {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, @@ -3644,7 +3758,9 @@ async def test_acount_pre_call_check_tokens_skips_without_max_input_tokens(monke monkeypatch.setattr(router, "get_router_model_info", lambda **kwargs: {}) calls = [] - monkeypatch.setattr(litellm, "token_counter", lambda *a, **k: calls.append(1) or 1000) + monkeypatch.setattr( + litellm, "token_counter", lambda *a, **k: calls.append(1) or 1000 + ) count = await router._acount_pre_call_check_tokens( model="m", @@ -3672,7 +3788,9 @@ def test_pre_call_checks_counts_tokens_from_responses_input_string(monkeypatch): ], enable_pre_call_checks=True, ) - monkeypatch.setattr(router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 1}) + monkeypatch.setattr( + router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 1} + ) deployments = [ {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, @@ -3697,7 +3815,9 @@ def test_pre_call_checks_counts_tokens_from_responses_input_list(monkeypatch): ], enable_pre_call_checks=True, ) - monkeypatch.setattr(router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 1}) + monkeypatch.setattr( + router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 1} + ) deployments = [ {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, @@ -3739,7 +3859,9 @@ def test_pre_call_checks_counts_responses_instructions_tokens(monkeypatch): ) assert with_instructions_tokens > input_only_tokens - monkeypatch.setattr(router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": input_only_tokens}) + monkeypatch.setattr( + router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": input_only_tokens} + ) with pytest.raises(litellm.ContextWindowExceededError): router._pre_call_checks( model="m", @@ -3808,7 +3930,9 @@ def test_pre_call_checks_counts_tool_definition_tokens(monkeypatch, prompt_kwarg prompt_only_tokens = router._count_pre_call_check_tokens( messages=prompt_kwargs.get("messages"), input=prompt_kwargs.get("input") ) - monkeypatch.setattr(router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": prompt_only_tokens}) + monkeypatch.setattr( + router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": prompt_only_tokens} + ) assert len(router._pre_call_checks(model="m", healthy_deployments=deployments, **prompt_kwargs)) == 1 with pytest.raises(litellm.ContextWindowExceededError): @@ -3928,7 +4052,7 @@ def test_count_pre_call_check_tokens_across_api_surfaces(): assert string_input_tokens > 0 assert list_input_tokens > 0 - with pytest.raises(ValueError, match="Either messages or input must be provided to count tokens"): + with pytest.raises(ValueError, match='Either messages or input must be provided to count tokens'): router._count_pre_call_check_tokens(messages=None, input=None) @@ -3943,7 +4067,9 @@ def test_pre_call_checks_no_messages_or_input_does_not_crash(monkeypatch): ], enable_pre_call_checks=True, ) - monkeypatch.setattr(router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 5}) + monkeypatch.setattr( + router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 5} + ) counted: list[dict] = [] original = router._count_pre_call_check_tokens @@ -4033,7 +4159,9 @@ def test_get_deployment_model_info_base_model_flow(): } # Test Case 1: Base model flow with custom model info that has base_model - with patch.object(litellm, "model_cost", {"test-custom-model": mock_custom_model_info}): + with patch.object( + litellm, "model_cost", {"test-custom-model": mock_custom_model_info} + ): with patch.object(litellm, "get_model_info") as mock_get_model_info: # Configure mock returns mock_get_model_info.side_effect = lambda model: { @@ -4041,11 +4169,15 @@ def test_get_deployment_model_info_base_model_flow(): "test-model": mock_litellm_model_name_info, }.get(model) - result = router.get_deployment_model_info(model_id="test-custom-model", model_name="test-model") + result = router.get_deployment_model_info( + model_id="test-custom-model", model_name="test-model" + ) # Verify that get_model_info was called for both base model and model name assert mock_get_model_info.call_count == 2 - mock_get_model_info.assert_any_call(model="gpt-3.5-turbo") # base model call + mock_get_model_info.assert_any_call( + model="gpt-3.5-turbo" + ) # base model call mock_get_model_info.assert_any_call(model="test-model") # model name call # Verify the result contains merged information @@ -4056,18 +4188,26 @@ def test_get_deployment_model_info_base_model_flow(): # 2. The result of step 1 gets merged into litellm_model_name_info (custom+base override litellm) # Fields from custom model (should override base model values) - assert result["input_cost_per_token"] == 0.001 # From custom model (overrides base 0.0015) - assert result["output_cost_per_token"] == 0.002 # From custom model (same as base) + assert ( + result["input_cost_per_token"] == 0.001 + ) # From custom model (overrides base 0.0015) + assert ( + result["output_cost_per_token"] == 0.002 + ) # From custom model (same as base) assert result["custom_field"] == "custom_value" # From custom model # Fields from base model that weren't overridden by custom assert result["max_tokens"] == 4096 # From base model assert result["litellm_provider"] == "openai" # From base model - assert result["mode"] == "chat" # From base model (overrides litellm "completion") + assert ( + result["mode"] == "chat" + ) # From base model (overrides litellm "completion") # The key field comes from base model since both base and litellm have it # and base model info overrides litellm model name info in final merge - assert result["key"] == "gpt-3.5-turbo" # From base model (overrides litellm key) + assert ( + result["key"] == "gpt-3.5-turbo" + ) # From base model (overrides litellm key) # Test Case 2: Custom model info without base_model mock_custom_model_info_no_base = { @@ -4086,7 +4226,9 @@ def test_get_deployment_model_info_base_model_flow(): "test-model": mock_litellm_model_name_info, }.get(model) - result = router.get_deployment_model_info(model_id="test-custom-model-no-base", model_name="test-model") + result = router.get_deployment_model_info( + model_id="test-custom-model-no-base", model_name="test-model" + ) # Should only call get_model_info once for model name (no base model) assert mock_get_model_info.call_count == 1 @@ -4106,7 +4248,9 @@ def test_get_deployment_model_info_base_model_flow(): "test-model": mock_litellm_model_name_info, }.get(model) - result = router.get_deployment_model_info(model_id="non-existent-model", model_name="test-model") + result = router.get_deployment_model_info( + model_id="non-existent-model", model_name="test-model" + ) # Should only call get_model_info once for model name assert mock_get_model_info.call_count == 1 @@ -4139,7 +4283,9 @@ def test_get_deployment_model_info_base_model_flow(): mock_get_model_info.side_effect = mock_get_model_info_side_effect - result = router.get_deployment_model_info(model_id="test-custom-model-invalid", model_name="test-model") + result = router.get_deployment_model_info( + model_id="test-custom-model-invalid", model_name="test-model" + ) # Should handle exception gracefully and still return merged result assert result is not None @@ -4148,8 +4294,12 @@ def test_get_deployment_model_info_base_model_flow(): # Test Case 5: Both model_cost.get() and get_model_info() return None with patch.object(litellm, "model_cost", {}): - with patch.object(litellm, "get_model_info", side_effect=Exception("Not found")): - result = router.get_deployment_model_info(model_id="non-existent", model_name="non-existent") + with patch.object( + litellm, "get_model_info", side_effect=Exception("Not found") + ): + result = router.get_deployment_model_info( + model_id="non-existent", model_name="non-existent" + ) # Should return None when no model info is found assert result is None @@ -4172,7 +4322,9 @@ def test_get_deployment_model_info_base_model_flow(): # Model NOT in built-in cost map — raise exception mock_get_model_info.side_effect = Exception("Model not in cost map") - result = router.get_deployment_model_info(model_id="custom-model-id", model_name="unknown-model") + result = router.get_deployment_model_info( + model_id="custom-model-id", model_name="unknown-model" + ) # Should return custom_model_info even when litellm_model_name_model_info is None assert result is not None @@ -4208,11 +4360,15 @@ def test_get_deployment_model_info_base_model_flow(): mock_get_model_info.side_effect = get_info_side_effect - result = router.get_deployment_model_info(model_id="custom-with-base", model_name="unknown-model") + result = router.get_deployment_model_info( + model_id="custom-with-base", model_name="unknown-model" + ) # Should return custom_model_info merged with base model info assert result is not None - assert result["input_cost_per_token"] == 0.01 # From custom (overrides base) + assert ( + result["input_cost_per_token"] == 0.01 + ) # From custom (overrides base) assert result["max_tokens"] == 8192 # From base model assert result["litellm_provider"] == "openai" # From base model @@ -4259,14 +4415,18 @@ def test_get_deployment_model_info_base_model_merge_priority(): "litellm_only_field": "litellm_value", } - with patch.object(litellm, "model_cost", {"custom-model-id": mock_custom_model_info}): + with patch.object( + litellm, "model_cost", {"custom-model-id": mock_custom_model_info} + ): with patch.object(litellm, "get_model_info") as mock_get_model_info: mock_get_model_info.side_effect = lambda model: { "gpt-4": mock_base_model_info, "test-model": mock_litellm_model_name_info, }.get(model) - result = router.get_deployment_model_info(model_id="custom-model-id", model_name="test-model") + result = router.get_deployment_model_info( + model_id="custom-model-id", model_name="test-model" + ) assert result is not None @@ -4276,17 +4436,29 @@ def test_get_deployment_model_info_base_model_merge_priority(): # 3. Result from steps 1-2 overrides litellm_model_name_info # Fields that should come from custom model info (highest priority) - assert result["input_cost_per_token"] == 0.01 # From custom model (overrides base 0.03) - assert result["max_tokens"] == 8000 # From custom model (overrides base 4096) + assert ( + result["input_cost_per_token"] == 0.01 + ) # From custom model (overrides base 0.03) + assert ( + result["max_tokens"] == 8000 + ) # From custom model (overrides base 4096) assert result["custom_only_field"] == "custom_value" # From custom model # Fields that should come from base model (not overridden by custom) - assert result["output_cost_per_token"] == 0.06 # From base model (not in custom) - assert result["litellm_provider"] == "openai" # From base model (not in custom) - assert result["base_only_field"] == "base_value" # From base model (not in custom) + assert ( + result["output_cost_per_token"] == 0.06 + ) # From base model (not in custom) + assert ( + result["litellm_provider"] == "openai" + ) # From base model (not in custom) + assert ( + result["base_only_field"] == "base_value" + ) # From base model (not in custom) # Fields that should come from litellm model name info (not overridden by custom+base) - assert result["mode"] == "completion" # From litellm model name info (not in custom or base) + assert ( + result["mode"] == "completion" + ) # From litellm model name info (not in custom or base) assert ( result["litellm_only_field"] == "litellm_value" ) # From litellm model name info (not in custom or base) @@ -4323,9 +4495,10 @@ def test_add_deployment_model_to_endpoint_for_llm_passthrough_route(): model="special-bedrock-model", model_name="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", ) - assert result["endpoint"] == "/model/us.anthropic.claude-haiku-4-5-20251001-v1:0/invoke", ( - f"Expected '/model/us.anthropic.claude-haiku-4-5-20251001-v1:0/invoke', got '{result['endpoint']}'" - ) + assert ( + result["endpoint"] + == "/model/us.anthropic.claude-haiku-4-5-20251001-v1:0/invoke" + ), f"Expected '/model/us.anthropic.claude-haiku-4-5-20251001-v1:0/invoke', got '{result['endpoint']}'" # Test Case 2: Bedrock invoke-with-response-stream endpoint kwargs = { @@ -4337,9 +4510,10 @@ def test_add_deployment_model_to_endpoint_for_llm_passthrough_route(): model="special-bedrock-model", model_name="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", ) - assert result["endpoint"] == "/model/us.anthropic.claude-haiku-4-5-20251001-v1:0/invoke-with-response-stream", ( - f"Expected streaming endpoint with stripped prefix, got '{result['endpoint']}'" - ) + assert ( + result["endpoint"] + == "/model/us.anthropic.claude-haiku-4-5-20251001-v1:0/invoke-with-response-stream" + ), f"Expected streaming endpoint with stripped prefix, got '{result['endpoint']}'" # Test Case 3: Bedrock converse endpoint kwargs = { @@ -4351,9 +4525,9 @@ def test_add_deployment_model_to_endpoint_for_llm_passthrough_route(): model="bedrock-model", model_name="bedrock/us.meta.llama3-8b-instruct-v1:0", ) - assert result["endpoint"] == "/model/us.meta.llama3-8b-instruct-v1:0/converse", ( - f"Expected '/model/us.meta.llama3-8b-instruct-v1:0/converse', got '{result['endpoint']}'" - ) + assert ( + result["endpoint"] == "/model/us.meta.llama3-8b-instruct-v1:0/converse" + ), f"Expected '/model/us.meta.llama3-8b-instruct-v1:0/converse', got '{result['endpoint']}'" # Test Case 4: Bedrock provider prefix auto-detected from model_name kwargs = { @@ -4364,9 +4538,9 @@ def test_add_deployment_model_to_endpoint_for_llm_passthrough_route(): model="router-model", model_name="bedrock/us.meta.llama3-8b-instruct-v1:0", ) - assert result["endpoint"] == "/model/us.meta.llama3-8b-instruct-v1:0/invoke", ( - f"Expected '/model/us.meta.llama3-8b-instruct-v1:0/invoke', got '{result['endpoint']}'" - ) + assert ( + result["endpoint"] == "/model/us.meta.llama3-8b-instruct-v1:0/invoke" + ), f"Expected '/model/us.meta.llama3-8b-instruct-v1:0/invoke', got '{result['endpoint']}'" def test_update_kwargs_with_deployment_uses_pass_through_request_timeout(): @@ -4418,10 +4592,14 @@ async def test_router_acompletion_with_unknown_model_and_default_fallback(): # Initialize the router with a default fallback router = litellm.Router(model_list=model_list, default_fallbacks=["gpt-4o"]) - messages = [{"role": "user", "content": "This call should succeed by falling back."}] + messages = [ + {"role": "user", "content": "This call should succeed by falling back."} + ] # Call completion with a model name that is NOT in the model_list - response = await router.acompletion(model="completely-unknown-model", messages=messages) + response = await router.acompletion( + model="completely-unknown-model", messages=messages + ) # Check that the call did not fail and we received a valid response object. assert response is not None @@ -4513,10 +4691,15 @@ def test_get_deployment_credentials_with_provider_aws_bedrock_runtime_endpoint() ], ) - credentials = router.get_deployment_credentials_with_provider(model_id="bedrock-claude-model") + credentials = router.get_deployment_credentials_with_provider( + model_id="bedrock-claude-model" + ) assert credentials is not None - assert credentials["aws_bedrock_runtime_endpoint"] == "https://bedrock-runtime.us-east-1.amazonaws.com" + assert ( + credentials["aws_bedrock_runtime_endpoint"] + == "https://bedrock-runtime.us-east-1.amazonaws.com" + ) assert credentials["aws_access_key_id"] == "test-access-key" assert credentials["aws_secret_access_key"] == "test-secret-key" assert credentials["aws_region_name"] == "us-east-1" @@ -4543,7 +4726,9 @@ def test_get_deployment_credentials_with_provider_includes_bucket_name(): ], ) - credentials = router.get_deployment_credentials_with_provider(model_id="vertex-gemini") + credentials = router.get_deployment_credentials_with_provider( + model_id="vertex-gemini" + ) assert credentials is not None assert credentials["gcs_bucket_name"] == "my-batch-bucket" @@ -4583,7 +4768,9 @@ def test_get_deployment_credentials_with_provider_resolves_credential_name(): ], ) - credentials = router.get_deployment_credentials_with_provider(model_id="azure-gpt-4") + credentials = router.get_deployment_credentials_with_provider( + model_id="azure-gpt-4" + ) assert credentials is not None assert credentials["api_key"] == "resolved-api-key" @@ -4619,7 +4806,9 @@ def test_get_deployment_credentials_with_provider_bedrock_batch_fields(): ], ) - credentials = router.get_deployment_credentials_with_provider(model_id="bedrock-batch-model") + credentials = router.get_deployment_credentials_with_provider( + model_id="bedrock-batch-model" + ) assert credentials is not None assert credentials["custom_llm_provider"] == "bedrock" @@ -4662,7 +4851,9 @@ def test_get_deployment_credentials_with_provider_preserves_aws_auth_params(): ], ) - credentials = router.get_deployment_credentials_with_provider(model_id="bedrock-batch-model") + credentials = router.get_deployment_credentials_with_provider( + model_id="bedrock-batch-model" + ) assert credentials is not None for key, value in aws_auth_params.items(): @@ -4697,11 +4888,15 @@ def test_get_deployment_credentials_with_provider_team_wildcard_priority(): ], ) - team_credentials = router.get_deployment_credentials_with_provider(model_id="openai/gpt-5.2", team_id="team-1") + team_credentials = router.get_deployment_credentials_with_provider( + model_id="openai/gpt-5.2", team_id="team-1" + ) assert team_credentials is not None assert team_credentials["api_key"] == "team-key" - global_credentials = router.get_deployment_credentials_with_provider(model_id="openai/gpt-5.2") + global_credentials = router.get_deployment_credentials_with_provider( + model_id="openai/gpt-5.2" + ) assert global_credentials is not None assert global_credentials["api_key"] == "global-key" @@ -4742,11 +4937,15 @@ def test_get_deployment_credentials_with_provider_skips_other_team_deployment(): assert other_team_credentials is not None assert other_team_credentials["vertex_project"] == "shared-project" - unscoped_credentials = router.get_deployment_credentials_with_provider(model_id="gemini-2.5-pro") + unscoped_credentials = router.get_deployment_credentials_with_provider( + model_id="gemini-2.5-pro" + ) assert unscoped_credentials is not None assert unscoped_credentials["vertex_project"] == "shared-project" - owner_credentials = router.get_deployment_credentials_with_provider(model_id="gemini-2.5-pro", team_id="team-b") + owner_credentials = router.get_deployment_credentials_with_provider( + model_id="gemini-2.5-pro", team_id="team-b" + ) assert owner_credentials is not None assert owner_credentials["vertex_project"] == "team-b-project" @@ -4773,8 +4972,16 @@ def test_get_deployment_credentials_with_provider_no_fallback_to_other_team_only ], ) - assert router.get_deployment_credentials_with_provider(model_id="gemini-2.5-pro", team_id="team-a") is None - assert router.get_deployment_credentials_with_provider(model_id="gemini-2.5-pro") is None + assert ( + router.get_deployment_credentials_with_provider( + model_id="gemini-2.5-pro", team_id="team-a" + ) + is None + ) + assert ( + router.get_deployment_credentials_with_provider(model_id="gemini-2.5-pro") + is None + ) def test_deployment_usable_by_team_helpers(): @@ -4814,7 +5021,9 @@ def test_deployment_usable_by_team_helpers(): assert router._deployment_usable_by_team(shared, "team-a") is True assert router._deployment_usable_by_team(shared, None) is True - picked = router._get_model_group_deployment_usable_by_team(model_group_name="gemini-2.5-pro", team_id="team-a") + picked = router._get_model_group_deployment_usable_by_team( + model_group_name="gemini-2.5-pro", team_id="team-a" + ) assert picked is not None assert picked.litellm_params.vertex_project == "shared-project" @@ -4824,7 +5033,12 @@ def test_deployment_usable_by_team_helpers(): assert owner_picked is not None assert owner_picked.litellm_params.vertex_project == "team-b-project" - assert router._get_model_group_deployment_usable_by_team(model_group_name="unknown-model", team_id="team-a") is None + assert ( + router._get_model_group_deployment_usable_by_team( + model_group_name="unknown-model", team_id="team-a" + ) + is None + ) def test_get_deployment_credentials_with_provider_skips_other_team_wildcard(): @@ -4856,7 +5070,9 @@ def test_get_deployment_credentials_with_provider_skips_other_team_wildcard(): assert other_team_credentials is not None assert other_team_credentials["api_key"] == "global-key" - owner_credentials = router.get_deployment_credentials_with_provider(model_id="openai/gpt-5.2", team_id="team-b") + owner_credentials = router.get_deployment_credentials_with_provider( + model_id="openai/gpt-5.2", team_id="team-b" + ) assert owner_credentials is not None assert owner_credentials["api_key"] == "team-b-key" @@ -4868,11 +5084,21 @@ def test_team_wildcard_credentials_not_usable_after_delete_deployment(): """ router = litellm.Router(model_list=[_team_wildcard_model(api_key="old-key")]) - assert router.get_deployment_credentials_with_provider(model_id="openai/gpt-5.2", team_id="team-1") is not None + assert ( + router.get_deployment_credentials_with_provider( + model_id="openai/gpt-5.2", team_id="team-1" + ) + is not None + ) router.delete_deployment(id="team-wildcard-id") - assert router.get_deployment_credentials_with_provider(model_id="openai/gpt-5.2", team_id="team-1") is None + assert ( + router.get_deployment_credentials_with_provider( + model_id="openai/gpt-5.2", team_id="team-1" + ) + is None + ) def test_global_wildcard_pattern_router_evicts_stale_entry_on_upsert_and_delete(): @@ -4945,13 +5171,22 @@ def test_team_wildcard_credentials_refreshed_on_upsert_and_set_model_list(): router = litellm.Router(model_list=[_team_wildcard_model(api_key="old-key")]) - router.upsert_deployment(deployment=Deployment(**_team_wildcard_model(api_key="new-key"))) - credentials = router.get_deployment_credentials_with_provider(model_id="openai/gpt-5.2", team_id="team-1") + router.upsert_deployment( + deployment=Deployment(**_team_wildcard_model(api_key="new-key")) + ) + credentials = router.get_deployment_credentials_with_provider( + model_id="openai/gpt-5.2", team_id="team-1" + ) assert credentials is not None assert credentials["api_key"] == "new-key" router.set_model_list(model_list=[]) - assert router.get_deployment_credentials_with_provider(model_id="openai/gpt-5.2", team_id="team-1") is None + assert ( + router.get_deployment_credentials_with_provider( + model_id="openai/gpt-5.2", team_id="team-1" + ) + is None + ) def test_get_available_guardrail_single_deployment(): @@ -5140,7 +5375,9 @@ async def test_anthropic_messages_call_type_is_cached(): startTime=1234567890.0, endTime=1234567891.0, completionStartTime=1234567890.5, - model_map_information=StandardLoggingModelInformation(model_map_key="gpt-3.5-turbo", model_map_value=None), + model_map_information=StandardLoggingModelInformation( + model_map_key="gpt-3.5-turbo", model_map_value=None + ), model="gpt-3.5-turbo", model_id="model-123", model_group="openai-gpt", @@ -5219,8 +5456,12 @@ async def test_anthropic_messages_call_type_is_cached(): ) # This assertion will FAIL if anthropic_messages is filtered out - assert cached_result is not None, "Model ID should be cached for anthropic_messages call type" - assert cached_result["model_id"] == test_model_id, f"Expected {test_model_id}, got {cached_result['model_id']}" + assert ( + cached_result is not None + ), "Model ID should be cached for anthropic_messages call type" + assert ( + cached_result["model_id"] == test_model_id + ), f"Expected {test_model_id}, got {cached_result['model_id']}" def test_update_kwargs_with_deployment_propagates_model_tags(): @@ -5245,7 +5486,9 @@ def test_update_kwargs_with_deployment_propagates_model_tags(): ) kwargs: dict = {"metadata": {}} - deployment = router.get_deployment_by_model_group_name(model_group_name="gpt-4o-mini") + deployment = router.get_deployment_by_model_group_name( + model_group_name="gpt-4o-mini" + ) router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) # Deployment tags should be propagated to kwargs metadata @@ -5274,7 +5517,9 @@ def test_update_kwargs_with_deployment_merges_tags_without_duplicates(): # Simulate request that already has tags (from request body or key/team level) kwargs: dict = {"metadata": {"tags": ["user-tag", "shared-tag"]}} - deployment = router.get_deployment_by_model_group_name(model_group_name="gpt-4o-mini") + deployment = router.get_deployment_by_model_group_name( + model_group_name="gpt-4o-mini" + ) router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) # Both sources should be merged, no duplicates @@ -5301,7 +5546,9 @@ def test_update_kwargs_with_deployment_no_tags(): ) kwargs: dict = {"metadata": {}} - deployment = router.get_deployment_by_model_group_name(model_group_name="gpt-4o-mini") + deployment = router.get_deployment_by_model_group_name( + model_group_name="gpt-4o-mini" + ) router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) # No tags key should be added if deployment has no tags @@ -5339,7 +5586,9 @@ def test_update_kwargs_with_deployment_merges_tools(): }, ], } - deployment = router.get_deployment_by_model_group_name(model_group_name="o3-deep-research") + deployment = router.get_deployment_by_model_group_name( + model_group_name="o3-deep-research" + ) router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) # Tools should be merged: deployment first, then request @@ -5370,7 +5619,9 @@ def test_update_kwargs_with_deployment_merge_tools_deployment_only(): ) kwargs: dict = {"metadata": {}} - deployment = router.get_deployment_by_model_group_name(model_group_name="o3-deep-research") + deployment = router.get_deployment_by_model_group_name( + model_group_name="o3-deep-research" + ) router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) assert kwargs["tools"] == [{"type": "web_search"}] @@ -5399,7 +5650,9 @@ def test_update_kwargs_with_deployment_merge_tools_request_overrides_tool_choice "metadata": {}, "tool_choice": "none", } - deployment = router.get_deployment_by_model_group_name(model_group_name="o3-deep-research") + deployment = router.get_deployment_by_model_group_name( + model_group_name="o3-deep-research" + ) router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) # Request tool_choice should be preserved (merged tools still applied) @@ -5501,8 +5754,12 @@ def test_update_kwargs_with_deployment_model_info_in_litellm_metadata(): ) kwargs: dict = {} - deployment = router.get_deployment_by_model_group_name(model_group_name="claude-sonnet-4") - router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs, function_name="generic_api_call") + deployment = router.get_deployment_by_model_group_name( + model_group_name="claude-sonnet-4" + ) + router._update_kwargs_with_deployment( + deployment=deployment, kwargs=kwargs, function_name="generic_api_call" + ) assert "litellm_metadata" in kwargs model_info = kwargs["litellm_metadata"]["model_info"] @@ -5534,8 +5791,12 @@ def test_update_kwargs_with_deployment_model_info_in_metadata(): ) kwargs: dict = {} - deployment = router.get_deployment_by_model_group_name(model_group_name="claude-sonnet-4") - router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs, function_name=None) + deployment = router.get_deployment_by_model_group_name( + model_group_name="claude-sonnet-4" + ) + router._update_kwargs_with_deployment( + deployment=deployment, kwargs=kwargs, function_name=None + ) assert "metadata" in kwargs model_info = kwargs["metadata"]["model_info"] @@ -5648,7 +5909,6 @@ async def test_acompletion_streaming_iterator_does_not_log_success_on_terminal_f initial_kwargs=dict(initial_kwargs), ) collected = [] - async def _drain(): async for chunk in result: collected.append(chunk) @@ -5675,7 +5935,6 @@ async def test_acompletion_streaming_iterator_does_not_log_success_on_terminal_f initial_kwargs=dict(initial_kwargs), ) collected = [] - async def _drain(): async for chunk in result: collected.append(chunk) @@ -5892,17 +6151,23 @@ def test_multiregion_team_deployments_unique_model_names(): assert len(deployments) == 0 # With team_id: O(n) scan finds BOTH regional deployments - deployments = router._get_all_deployments(model_name="claude-sonnet", team_id="metis-team") + deployments = router._get_all_deployments( + model_name="claude-sonnet", team_id="metis-team" + ) assert len(deployments) == 2 deployment_names = {d["model_name"] for d in deployments} assert deployment_names == {"metis-claude-us-east-1", "metis-claude-us-west-2"} # Each deployment has a unique ID (critical for cooldown/retry to work) deployment_ids = {d["model_info"]["id"] for d in deployments} - assert len(deployment_ids) == 2, "Each deployment must have a unique ID for cooldown tracking" + assert ( + len(deployment_ids) == 2 + ), "Each deployment must have a unique ID for cooldown tracking" # Wrong team: returns nothing - deployments = router._get_all_deployments(model_name="claude-sonnet", team_id="other-team") + deployments = router._get_all_deployments( + model_name="claude-sonnet", team_id="other-team" + ) assert len(deployments) == 0 @@ -5947,8 +6212,12 @@ async def test_multiregion_team_failover_between_regions(): ) # Verify the router finds both deployments for the team - deployments = router._get_all_deployments(model_name="claude-sonnet", team_id="metis-team") - assert len(deployments) == 2, "Router must find both regional deployments by team_public_model_name" + deployments = router._get_all_deployments( + model_name="claude-sonnet", team_id="metis-team" + ) + assert ( + len(deployments) == 2 + ), "Router must find both regional deployments by team_public_model_name" # Make a normal request — should succeed from one of the regions response = await router.acompletion( @@ -6073,7 +6342,9 @@ def test_explicit_model_access_does_not_force_access_group_filtering(): }, ) - deployment_groups = [d.get("model_info", {}).get("access_groups") for d in deployments] + deployment_groups = [ + d.get("model_info", {}).get("access_groups") for d in deployments + ] assert ["AG1"] in deployment_groups assert ["AG2"] in deployment_groups @@ -6118,7 +6389,9 @@ def test_access_group_filter_empty_does_not_bypass_via_litellm_model_fallback( orig_groups = router.get_model_access_groups - def fake_get_model_access_groups(model_name=None, model_access_group=None, team_id=None): + def fake_get_model_access_groups( + model_name=None, model_access_group=None, team_id=None + ): if model_name == "gpt-5" and model_access_group is None: return {"AG1": ["gpt-5"], "AG2": ["gpt-5"]} return orig_groups( @@ -6193,7 +6466,9 @@ def test_access_group_block_does_not_silently_use_default_fallback_model( orig_groups = router.get_model_access_groups - def fake_get_model_access_groups(model_name=None, model_access_group=None, team_id=None): + def fake_get_model_access_groups( + model_name=None, model_access_group=None, team_id=None + ): if model_name == "gpt-5" and model_access_group is None: return {"AG1": ["gpt-5"], "AG2": ["gpt-5"]} return orig_groups( @@ -6260,7 +6535,9 @@ def test_access_group_block_via_litellm_model_branch_does_not_use_default_fallba orig_groups = router.get_model_access_groups - def fake_get_model_access_groups(model_name=None, model_access_group=None, team_id=None): + def fake_get_model_access_groups( + model_name=None, model_access_group=None, team_id=None + ): if model_name == "gpt-5" and model_access_group is None: return {"AG1": ["gpt-5"], "AG2": ["gpt-5"]} return orig_groups( @@ -6315,7 +6592,9 @@ def test_try_early_resolve_deployments_for_model_not_in_names(): ) assert ( - router_in_names._try_early_resolve_deployments_for_model_not_in_names(model="gpt-5", request_team_id=None) + router_in_names._try_early_resolve_deployments_for_model_not_in_names( + model="gpt-5", request_team_id=None + ) is None ) assert ( @@ -6337,8 +6616,10 @@ def test_try_early_resolve_deployments_for_model_not_in_names(): ] ) - pattern_result = pattern_router._try_early_resolve_deployments_for_model_not_in_names( - model="openai/gpt-4o-mini", request_team_id=None + pattern_result = ( + pattern_router._try_early_resolve_deployments_for_model_not_in_names( + model="openai/gpt-4o-mini", request_team_id=None + ) ) assert pattern_result is not None resolved_model, pattern_deployments = pattern_result @@ -6364,8 +6645,10 @@ def test_try_early_resolve_deployments_for_model_not_in_names(): }, } - default_result = default_router._try_early_resolve_deployments_for_model_not_in_names( - model="brand-new-model", request_team_id=None + default_result = ( + default_router._try_early_resolve_deployments_for_model_not_in_names( + model="brand-new-model", request_team_id=None + ) ) assert default_result is not None resolved_model, default_deployment = default_result @@ -6373,7 +6656,10 @@ def test_try_early_resolve_deployments_for_model_not_in_names(): assert isinstance(default_deployment, dict) assert default_deployment["litellm_params"]["model"] == "brand-new-model" # The original default_deployment must not be mutated. - assert default_router.default_deployment["litellm_params"]["model"] == "openai/will-be-overridden" + assert ( + default_router.default_deployment["litellm_params"]["model"] + == "openai/will-be-overridden" + ) def _router_with_two_deployments(blocked_flags): @@ -6421,7 +6707,10 @@ def _seed_unhealthy_states(router, unhealthy_ids, timestamp=None): ts = timestamp if timestamp is not None else time.time() router.health_state_cache.set_deployment_health_states( - {uid: {"is_healthy": False, "timestamp": ts, "reason": "test_unhealthy"} for uid in unhealthy_ids} + { + uid: {"is_healthy": False, "timestamp": ts, "reason": "test_unhealthy"} + for uid in unhealthy_ids + } ) @@ -6492,7 +6781,9 @@ async def test_async_get_fully_unhealthy_model_names_noop_with_allowed_fails_pol @pytest.mark.asyncio async def test_async_get_healthy_deployments_skips_blocked_deployment(): router = _router_with_two_deployments([True, False]) - healthy, all_dep = await router._async_get_healthy_deployments(model="gpt-4o", parent_otel_span=None) + healthy, all_dep = await router._async_get_healthy_deployments( + model="gpt-4o", parent_otel_span=None + ) healthy_ids = [d["model_info"]["id"] for d in healthy] assert "dep-0" not in healthy_ids assert "dep-1" in healthy_ids @@ -6501,7 +6792,9 @@ async def test_async_get_healthy_deployments_skips_blocked_deployment(): def test_get_healthy_deployments_sync_skips_blocked_deployment(): router = _router_with_two_deployments([False, True]) - healthy, all_dep = router._get_healthy_deployments(model="gpt-4o", parent_otel_span=None) + healthy, all_dep = router._get_healthy_deployments( + model="gpt-4o", parent_otel_span=None + ) healthy_ids = [d["model_info"]["id"] for d in healthy] assert "dep-0" in healthy_ids assert "dep-1" not in healthy_ids @@ -6518,7 +6811,9 @@ def test_filter_blocked_deployments_drops_blocked_keeps_unblocked(): @pytest.mark.asyncio async def test_public_async_get_healthy_deployments_skips_blocked_on_primary_path(): router = _router_with_two_deployments([True, False]) - deployments = await router.async_get_healthy_deployments(model="gpt-4o", request_kwargs={}) + deployments = await router.async_get_healthy_deployments( + model="gpt-4o", request_kwargs={} + ) assert isinstance(deployments, list) ids = [d["model_info"]["id"] for d in deployments] assert "dep-0" not in ids @@ -6560,7 +6855,9 @@ def _router_with_two_pass_through_deployments(blocked_flags): def test_get_available_deployment_for_pass_through_skips_blocked(): router = _router_with_two_pass_through_deployments([True, False]) - deployment = router.get_available_deployment_for_pass_through(model="gpt-4o", request_kwargs={}) + deployment = router.get_available_deployment_for_pass_through( + model="gpt-4o", request_kwargs={} + ) assert deployment["model_info"]["id"] == "pt-1" @@ -6569,7 +6866,9 @@ def test_get_available_deployment_for_pass_through_raises_when_dict_blocked(): router = _router_with_two_pass_through_deployments([True, True]) with pytest.raises(litellm.ServiceUnavailableError): - router.get_available_deployment_for_pass_through(model="pt-0", request_kwargs={}) + router.get_available_deployment_for_pass_through( + model="pt-0", request_kwargs={} + ) def test_initialize_deployment_for_pass_through_keeps_bedrock_iam_deployment(): @@ -6593,7 +6892,9 @@ def test_initialize_deployment_for_pass_through_keeps_bedrock_iam_deployment(): } ] ) - assert [m["model_info"]["id"] for m in router.get_model_list()] == ["bedrock-iam-pt"] + assert [m["model_info"]["id"] for m in router.get_model_list()] == [ + "bedrock-iam-pt" + ] def test_pass_through_deployment_api_key_resolves_via_get_credentials(): @@ -6604,7 +6905,12 @@ def test_pass_through_deployment_api_key_resolves_via_get_credentials(): router = _router_with_two_pass_through_deployments([False, False]) passthrough_router = PassthroughEndpointRouter(llm_router_getter=lambda: router) assert len(router.get_model_list()) == 2 - assert passthrough_router.get_credentials(custom_llm_provider="openai", region_name=None) == "sk-fake-for-tests" + assert ( + passthrough_router.get_credentials( + custom_llm_provider="openai", region_name=None + ) + == "sk-fake-for-tests" + ) def test_get_deployment_credentials_returns_none_for_blocked_deployment(): @@ -6638,9 +6944,16 @@ def test_is_deployment_blocked_static_helper_reflects_blocked_flag(): # No model_info on deployment object → treated as not blocked assert litellm.Router._is_deployment_blocked(object()) is False missing_blocked = types.SimpleNamespace() - assert litellm.Router._is_deployment_blocked(types.SimpleNamespace(model_info=missing_blocked)) is False assert ( - litellm.Router._is_deployment_blocked(types.SimpleNamespace(model_info=types.SimpleNamespace(blocked=True))) + litellm.Router._is_deployment_blocked( + types.SimpleNamespace(model_info=missing_blocked) + ) + is False + ) + assert ( + litellm.Router._is_deployment_blocked( + types.SimpleNamespace(model_info=types.SimpleNamespace(blocked=True)) + ) is True ) @@ -6680,7 +6993,9 @@ class TestRouterRequestTimeoutPropagation: litellm.request_timeout = original_value litellm.request_timeout_explicitly_set = original_flag - def test_request_timeout_stored_independently_when_both_set(self, explicit_request_timeout): + def test_request_timeout_stored_independently_when_both_set( + self, explicit_request_timeout + ): router = self._make_router(timeout=330) assert router.timeout == 330 assert router.request_timeout == 300 @@ -6698,16 +7013,22 @@ class TestRouterRequestTimeoutPropagation: litellm.request_timeout = original_value litellm.request_timeout_explicitly_set = original_flag - def test_non_stream_prefers_request_timeout_over_router_timeout(self, explicit_request_timeout): + def test_non_stream_prefers_request_timeout_over_router_timeout( + self, explicit_request_timeout + ): router = self._make_router(timeout=330) assert router._get_non_stream_timeout(kwargs={}, data={}) == 300 - def test_stream_prefers_request_timeout_over_router_timeout(self, explicit_request_timeout): + def test_stream_prefers_request_timeout_over_router_timeout( + self, explicit_request_timeout + ): router = self._make_router(timeout=330) # stream=True resolves through _get_stream_timeout; request_timeout must win. assert router._get_timeout(kwargs={"stream": True}, data={}) == 300 - def test_explicit_stream_timeout_still_wins_over_request_timeout(self, explicit_request_timeout): + def test_explicit_stream_timeout_still_wins_over_request_timeout( + self, explicit_request_timeout + ): router = self._make_router(timeout=330, stream_timeout=45) assert router._get_stream_timeout(kwargs={}, data={}) == 45 @@ -6723,13 +7044,22 @@ class TestRouterRequestTimeoutPropagation: litellm.request_timeout = original_value litellm.request_timeout_explicitly_set = original_flag - def test_per_deployment_timeout_overrides_request_timeout(self, explicit_request_timeout): + def test_per_deployment_timeout_overrides_request_timeout( + self, explicit_request_timeout + ): router = self._make_router(timeout=330) assert router._get_non_stream_timeout(kwargs={}, data={"timeout": 120}) == 120 - def test_per_request_timeout_overrides_request_timeout(self, explicit_request_timeout): + def test_per_request_timeout_overrides_request_timeout( + self, explicit_request_timeout + ): router = self._make_router(timeout=330) - assert router._get_non_stream_timeout(kwargs={"timeout": 60}, data={"timeout": 120}) == 60 + assert ( + router._get_non_stream_timeout( + kwargs={"timeout": 60}, data={"timeout": 120} + ) + == 60 + ) # --------------------------------------------------------------------------- @@ -7038,7 +7368,9 @@ class TestAdvisorSubCallCooldown: ) def _cooled_down_ids(self, router): - active = router.cooldown_cache.get_active_cooldowns(model_ids=["dep-1"], parent_otel_span=None) + active = router.cooldown_cache.get_active_cooldowns( + model_ids=["dep-1"], parent_otel_span=None + ) return [entry[0] for entry in active] @pytest.mark.asyncio @@ -7047,7 +7379,12 @@ class TestAdvisorSubCallCooldown: router = self._router() now = datetime.now() - assert router.deployment_callback_on_failure(self._kwargs(self._auth_error()), None, now, now) is True + assert ( + router.deployment_callback_on_failure( + self._kwargs(self._auth_error()), None, now, now + ) + is True + ) assert "dep-1" in self._cooled_down_ids(router) def test_advisor_orchestration_failure_does_not_cool_down_deployment(self): @@ -7062,7 +7399,12 @@ class TestAdvisorSubCallCooldown: mark_advisor_orchestration_failure(exception) now = datetime.now() - assert router.deployment_callback_on_failure(self._kwargs(exception), None, now, now) is False + assert ( + router.deployment_callback_on_failure( + self._kwargs(exception), None, now, now + ) + is False + ) assert "dep-1" not in self._cooled_down_ids(router) @@ -7119,13 +7461,13 @@ def test_stream_chunks_have_generated_content_detects_text_and_non_text(): audio_chunk = _chunk(audio_delta) assert _stream_chunks_have_generated_content([audio_chunk]) is True - images_delta = Delta( - images=[{"image_url": {"url": "https://example.com/img.png"}, "index": 0, "type": "image_url"}] - ) + images_delta = Delta(images=[{"image_url": {"url": "https://example.com/img.png"}, "index": 0, "type": "image_url"}]) images_chunk = _chunk(images_delta) assert _stream_chunks_have_generated_content([images_chunk]) is True - annotations_delta = Delta(annotations=[{"type": "url_citation", "url_citation": {"url": "https://example.com"}}]) + annotations_delta = Delta( + annotations=[{"type": "url_citation", "url_citation": {"url": "https://example.com"}}] + ) annotations_chunk = _chunk(annotations_delta) assert _stream_chunks_have_generated_content([annotations_chunk]) is True @@ -7169,8 +7511,12 @@ def test_get_configured_token_limits_skips_wildcard_pattern_matching(): ] ) - with patch.object(router.pattern_router, "route", side_effect=AssertionError("pattern route called")): - assert router.get_configured_token_limits("bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0") == (None, None) + with patch.object( + router.pattern_router, "route", side_effect=AssertionError("pattern route called") + ): + assert router.get_configured_token_limits( + "bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0" + ) == (None, None) def test_get_configured_token_limits_treats_malformed_values_as_absent(): @@ -7243,8 +7589,13 @@ def test_get_configured_mode_skips_wildcard_pattern_matching(): ] ) - with patch.object(router.pattern_router, "route", side_effect=AssertionError("pattern route called")): - assert router.get_configured_mode("bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0") is None + with patch.object( + router.pattern_router, "route", side_effect=AssertionError("pattern route called") + ): + assert ( + router.get_configured_mode("bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0") + is None + ) def test_get_configured_mode_treats_malformed_values_as_absent(): @@ -7303,8 +7654,13 @@ def test_get_configured_display_name_skips_wildcard_pattern_matching(): ] ) - with patch.object(router.pattern_router, "route", side_effect=AssertionError("pattern route called")): - assert router.get_configured_display_name("bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0") is None + with patch.object( + router.pattern_router, "route", side_effect=AssertionError("pattern route called") + ): + assert ( + router.get_configured_display_name("bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0") + is None + ) def test_get_configured_display_name_treats_malformed_values_as_absent(): @@ -7537,16 +7893,13 @@ async def test_acreate_batch_request_bedrock_tags_override_deployment_tags(): mock_client = MagicMock() mock_client.post = AsyncMock(side_effect=lambda *args, **kwargs: fake_response()) - with ( - patch.object( - CommonBatchFilesUtils, - "sign_aws_request", - return_value=({"Authorization": "signed"}, b"{}"), - ) as mock_sign, - patch( - "litellm.llms.custom_httpx.llm_http_handler.get_async_httpx_client", - return_value=mock_client, - ), + with patch.object( + CommonBatchFilesUtils, + "sign_aws_request", + return_value=({"Authorization": "signed"}, b"{}"), + ) as mock_sign, patch( + "litellm.llms.custom_httpx.llm_http_handler.get_async_httpx_client", + return_value=mock_client, ): await router.acreate_batch( model="bedrock-batch-model", @@ -7575,13 +7928,13 @@ async def test_avector_store_search_injects_router(): """ from litellm.types.vector_stores import VectorStoreSearchResponse - expected_response = VectorStoreSearchResponse(object="vector_store.search_results.page", search_query="q", data=[]) + expected_response = VectorStoreSearchResponse( + object="vector_store.search_results.page", search_query="q", data=[] + ) mock_asearch = AsyncMock(return_value=expected_response) # Router.__init__ binds asearch via a local import, so patch the module # attribute before constructing the Router. - with patch( - "litellm.vector_stores.main.asearch", new=mock_asearch - ): # test-quality-ok: the SDK call is the only place the injected router kwarg is observable + with patch("litellm.vector_stores.main.asearch", new=mock_asearch): # test-quality-ok: the SDK call is the only place the injected router kwarg is observable router = litellm.Router( model_list=[ { @@ -7615,9 +7968,7 @@ async def test_avector_store_create_does_not_inject_router(): } ] ) - with patch( - "litellm.vector_stores.main.acreate", new=mock_acreate - ): # test-quality-ok: the SDK call is the only place a leaked router kwarg would surface + with patch("litellm.vector_stores.main.acreate", new=mock_acreate): # test-quality-ok: the SDK call is the only place a leaked router kwarg would surface create_response = await router.avector_store_create(model=None, custom_llm_provider="openai") assert create_response is expected_response @@ -7633,13 +7984,13 @@ def test_vector_store_search_injects_router(): """ from litellm.types.vector_stores import VectorStoreSearchResponse - expected_response = VectorStoreSearchResponse(object="vector_store.search_results.page", search_query="q", data=[]) + expected_response = VectorStoreSearchResponse( + object="vector_store.search_results.page", search_query="q", data=[] + ) mock_search = MagicMock(return_value=expected_response) # Router.__init__ binds search via a local import, so patch the module # attribute before constructing the Router. - with patch( - "litellm.vector_stores.main.search", new=mock_search - ): # test-quality-ok: the SDK call is the only place the injected router kwarg is observable + with patch("litellm.vector_stores.main.search", new=mock_search): # test-quality-ok: the SDK call is the only place the injected router kwarg is observable router = litellm.Router( model_list=[ { @@ -7648,7 +7999,9 @@ def test_vector_store_search_injects_router(): } ] ) - search_response = router.vector_store_search(vector_store_id="v", query="q", custom_llm_provider="s3_vectors") + search_response = router.vector_store_search( + vector_store_id="v", query="q", custom_llm_provider="s3_vectors" + ) assert search_response is expected_response mock_search.assert_called_once() @@ -7662,9 +8015,7 @@ def test_vector_store_create_does_not_inject_router(): mock_create = MagicMock(return_value=expected_response) # Router.__init__ binds create via a local import, so patch the module # attribute before constructing the Router. - with patch( - "litellm.vector_stores.main.create", new=mock_create - ): # test-quality-ok: the SDK call is the only place a leaked router kwarg would surface + with patch("litellm.vector_stores.main.create", new=mock_create): # test-quality-ok: the SDK call is the only place a leaked router kwarg would surface router = litellm.Router( model_list=[ { @@ -7699,7 +8050,9 @@ class TestPreRoutingStrategyRegistryLifecycle: def _complexity_router_params(default_model: str, tags=None) -> dict: return { "model": "auto_router/complexity_router", - "complexity_router_config": {"tiers": {"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o", "COMPLEX": "gpt-4o"}}, + "complexity_router_config": { + "tiers": {"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o", "COMPLEX": "gpt-4o"} + }, "complexity_router_default_model": default_model, **({"tags": tags} if tags else {}), } @@ -8004,7 +8357,9 @@ class TestPreRoutingStrategyRegistryLifecycle: deployment=Deployment( model_name="hybrid-router", litellm_params=LiteLLM_Params( - **self._hybrid_router_params({"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o", "COMPLEX": "gpt-4o"}) + **self._hybrid_router_params( + {"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o", "COMPLEX": "gpt-4o"} + ) ), model_info=ModelInfo(id="router-1", db_model=True), ) @@ -8113,7 +8468,9 @@ class TestPreRoutingStrategyRegistryLifecycle: ({"model": "openai/gpt-4o"}, False), ] for params, expected in cases: - actual = router._deployment_participates_in_adaptive_routing(litellm_params=LiteLLM_Params(**params)) + actual = router._deployment_participates_in_adaptive_routing( + litellm_params=LiteLLM_Params(**params) + ) assert actual is expected, params["model"] @@ -8300,16 +8657,22 @@ class TestUpsertDeploymentRollback: router.delete_deployment(id="prod-1") assert router.has_model_id("prod-1") is False - router._restore_deployment_after_failed_upsert(previous_deployment=previous, model_id="prod-1") + router._restore_deployment_after_failed_upsert( + previous_deployment=previous, model_id="prod-1" + ) restored = router.get_deployment(model_id="prod-1") assert restored is not None assert restored.litellm_params.model == "gpt-4o" - router._restore_deployment_after_failed_upsert(previous_deployment=previous, model_id="prod-1") + router._restore_deployment_after_failed_upsert( + previous_deployment=previous, model_id="prod-1" + ) assert len(router.model_list) == 1 - router._restore_deployment_after_failed_upsert(previous_deployment=None, model_id="prod-1") + router._restore_deployment_after_failed_upsert( + previous_deployment=None, model_id="prod-1" + ) assert len(router.model_list) == 1 @@ -9075,14 +9438,18 @@ class TestAutoRouterSharedModelNameConnectionParams: return httpx.Response( status_code=200, json={ - "candidates": [{"content": {"parts": [{"text": "Paris"}], "role": "model"}, "finishReason": "STOP"}], + "candidates": [ + {"content": {"parts": [{"text": "Paris"}], "role": "model"}, "finishReason": "STOP"} + ], "usageMetadata": {"promptTokenCount": 5, "candidatesTokenCount": 1, "totalTokenCount": 6}, "modelVersion": "gemini-3.6-flash", }, request=httpx.Request("POST", "https://generativelanguage.googleapis.com"), ) - @pytest.mark.parametrize("plain_entry_first", [True, False], ids=["plain_entry_first", "marker_entry_first"]) + @pytest.mark.parametrize( + "plain_entry_first", [True, False], ids=["plain_entry_first", "marker_entry_first"] + ) async def test_routed_tier_call_goes_out_on_its_own_endpoint_and_credentials(self, plain_entry_first): """The outbound provider request for the routed tier hits the tier's own Gemini host with the tier's own key, never the plain sibling's api_base or api_key.""" @@ -9205,7 +9572,9 @@ async def _drive_cyclic_fallback(router, capture, recorder=None, **request_kwarg litellm.callbacks.append(recorder) try: with pytest.raises(litellm.InternalServerError): - await router.acompletion(model="group-a", messages=[{"role": "user", "content": "hi"}], **request_kwargs) + await router.acompletion( + model="group-a", messages=[{"role": "user", "content": "hi"}], **request_kwargs + ) finally: router_logger.removeHandler(capture) router_logger.setLevel(previous_level) @@ -9244,9 +9613,9 @@ async def test_retry_breadcrumbs_do_not_carry_the_walk_state(): breadcrumbs = [breadcrumb for hop in recorder.breadcrumbs_per_target for breadcrumb in hop] assert breadcrumbs, "no retry breadcrumbs were recorded" - assert any("fallback_depth" in breadcrumb for breadcrumb in breadcrumbs), ( - "no breadcrumb carried router walk state, so this test cannot see the leak" - ) + assert any( + "fallback_depth" in breadcrumb for breadcrumb in breadcrumbs + ), "no breadcrumb carried router walk state, so this test cannot see the leak" for breadcrumb in breadcrumbs: assert "attempted_targets" not in breadcrumb @@ -9292,9 +9661,7 @@ async def test_retry_breadcrumbs_never_carry_a_forwarded_credential(container_ke breadcrumbs = metadata["previous_models"] assert breadcrumbs, "no retry breadcrumbs were recorded" dumped = json.dumps(breadcrumbs, default=str) - assert container_key in dumped, ( - "the credential-bearing kwarg never reached the breadcrumb, so this test cannot see the leak" - ) + assert container_key in dumped, "the credential-bearing kwarg never reached the breadcrumb, so this test cannot see the leak" assert _BREADCRUMB_CREDENTIAL_CANARY not in dumped @@ -9399,7 +9766,9 @@ async def test_fallback_failure_detail_from_upstream_is_bounded(): await _drive_cyclic_fallback( _cyclic_fallback_router(), capture, - mock_response=litellm.InternalServerError(message=huge_message, llm_provider="openai", model="group-a"), + mock_response=litellm.InternalServerError( + message=huge_message, llm_provider="openai", model="group-a" + ), ) assert capture.messages, "the fallback failure path did not log at ERROR" @@ -9465,7 +9834,9 @@ def test_ensure_deployment_affinity_callback_is_idempotent(): try: router._ensure_deployment_affinity_callback() router._ensure_deployment_affinity_callback() - affinity_callbacks = [cb for cb in router.optional_callbacks or [] if isinstance(cb, DeploymentAffinityCheck)] + affinity_callbacks = [ + cb for cb in router.optional_callbacks or [] if isinstance(cb, DeploymentAffinityCheck) + ] assert len(affinity_callbacks) == 1 finally: for cb in router.optional_callbacks or []: @@ -9607,7 +9978,9 @@ class TestModelGroupAliasReachesPreRoutingStrategies: router = self._router("auto_routers") metadata: dict = {} - response = await router.acompletion(model="smart-alias", messages=self._messages(), metadata=metadata) + response = await router.acompletion( + model="smart-alias", messages=self._messages(), metadata=metadata + ) assert response.choices[0].message.content == "routed by the tier" assert metadata["model_group"] == "smart-alias" @@ -9829,6 +10202,8 @@ class TestAutoRouterCompressionDecoupling: 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 @@ -9855,14 +10230,17 @@ class TestAzureBaseModelFallbackLogging: def test_map_known_deployment_name_resolves_without_error_log(self): router = self._router_with_azure_deployment("azure/gpt-4o") - with patch("litellm.router.verbose_router_logger.error") as mock_error: + with patch( + "litellm.router.verbose_router_logger.error" + ) as mock_error: model_info = router.get_router_model_info( deployment=None, received_model_name="my-group", id="azure-base-model-test-id" ) - assert not any("Could not identify azure model" in str(call) for call in mock_error.call_args_list), ( - f"unexpected error log: {mock_error.call_args_list}" - ) + assert not any( + "Could not identify azure model" in str(call) + for call in mock_error.call_args_list + ), f"unexpected error log: {mock_error.call_args_list}" # the fallback resolution must actually surface the map values assert model_info["max_input_tokens"] == litellm.model_cost["azure/gpt-4o"]["max_input_tokens"] assert model_info["input_cost_per_token"] == litellm.model_cost["azure/gpt-4o"]["input_cost_per_token"] @@ -9870,14 +10248,17 @@ class TestAzureBaseModelFallbackLogging: def test_unmappable_deployment_name_still_logs_error(self): router = self._router_with_azure_deployment("azure/my-custom-deployment-name") - with patch("litellm.router.verbose_router_logger.error") as mock_error: + with patch( + "litellm.router.verbose_router_logger.error" + ) as mock_error: model_info = router.get_router_model_info( deployment=None, received_model_name="my-group", id="azure-base-model-test-id" ) - assert any("Could not identify azure model" in str(call) for call in mock_error.call_args_list), ( - "expected the error log for an unmappable azure deployment name" - ) + assert any( + "Could not identify azure model" in str(call) + for call in mock_error.call_args_list + ), "expected the error log for an unmappable azure deployment name" # unmappable names resolve to a zeroed stub — unchanged behavior assert model_info.get("max_input_tokens") is None @@ -9904,7 +10285,6 @@ class TestAzureBaseModelFallbackLogging: ) assert model_info["max_input_tokens"] == litellm.model_cost["azure/gpt-4o-mini"]["max_input_tokens"] - def test_model_group_info_intersects_supported_reasoning_efforts(): router = litellm.Router( model_list=[ @@ -9995,6 +10375,7 @@ def test_model_group_info_reasoning_efforts_are_unknown_when_any_deployment_is_o assert result.supported_reasoning_efforts is None + def test_model_group_info_surfaces_supports_parallel_function_calling(local_model_cost_map): """``/model_group/info`` folds each deployment's registry flags into the group; a deployment whose registry entry declares parallel function calling must flip the group to True instead of False.""" @@ -10227,7 +10608,6 @@ class TestAddDeploymentApiBaseProviderResolution: assert deployment is not None assert deployment.litellm_params.custom_llm_provider == "openai" - # ===================================================================== # anthropic_messages mid-stream-fallback helpers, added for #24004 # (mid-stream fallback not supported for anthropic_messages route type). @@ -10338,7 +10718,10 @@ class _AnthropicMessagesFallbackByteStream: def _anthropic_messages_overloaded_error_chunk() -> bytes: - return b'event: error\ndata: {"type": "error", "error": {"type": "overloaded_error", "message": "Overloaded"}}\n\n' + return ( + b"event: error\n" + b'data: {"type": "error", "error": {"type": "overloaded_error", "message": "Overloaded"}}\n\n' + ) def _anthropic_messages_invalid_request_error_chunk() -> bytes: @@ -10383,7 +10766,9 @@ async def test_anthropic_messages_streaming_iterator_passthrough(): [_anthropic_messages_content_chunk("hi"), _anthropic_messages_content_chunk(" there")] ) - wrapped = await router._aanthropic_messages_streaming_iterator(response=source, initial_kwargs={"model": "primary"}) + wrapped = await router._aanthropic_messages_streaming_iterator( + response=source, initial_kwargs={"model": "primary"} + ) collected = [chunk async for chunk in wrapped] assert collected == [_anthropic_messages_content_chunk("hi"), _anthropic_messages_content_chunk(" there")] @@ -10402,14 +10787,12 @@ async def test_anthropic_messages_streaming_iterator_flushes_buffered_lifecycle_ [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("hi"), message_stop] ) - wrapped = await router._aanthropic_messages_streaming_iterator(response=source, initial_kwargs={"model": "primary"}) + wrapped = await router._aanthropic_messages_streaming_iterator( + response=source, initial_kwargs={"model": "primary"} + ) collected = [chunk async for chunk in wrapped] - assert collected == [ - _anthropic_messages_message_start_chunk(), - _anthropic_messages_content_chunk("hi"), - message_stop, - ] + assert collected == [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("hi"), message_stop] @pytest.mark.asyncio @@ -10421,7 +10804,9 @@ async def test_anthropic_messages_streaming_iterator_flushes_buffered_frames_on_ message_stop = b'event: message_stop\ndata: {"type": "message_stop"}\n\n' source = _AnthropicMessagesFakeByteStream([_anthropic_messages_message_start_chunk(), message_stop]) - wrapped = await router._aanthropic_messages_streaming_iterator(response=source, initial_kwargs={"model": "primary"}) + wrapped = await router._aanthropic_messages_streaming_iterator( + response=source, initial_kwargs={"model": "primary"} + ) collected = [chunk async for chunk in wrapped] assert collected == [_anthropic_messages_message_start_chunk(), message_stop] @@ -10492,9 +10877,7 @@ async def test_anthropic_messages_leading_ping_keepalive_is_forwarded_live(): yield _anthropic_messages_message_start_chunk() yield _anthropic_messages_content_chunk("hi") - wrapped = await router._aanthropic_messages_streaming_iterator( - response=source(), initial_kwargs={"model": "primary"} - ) + wrapped = await router._aanthropic_messages_streaming_iterator(response=source(), initial_kwargs={"model": "primary"}) assert await asyncio.wait_for(wrapped.__anext__(), timeout=1) == _anthropic_messages_ping_chunk() content_released.set() @@ -12295,9 +12678,9 @@ class TestTierParamsTheTargetAccepts: def test_declared_param_allowlist_ignores_malformed_declarations(self): """A str is iterable, so without the type guard a YAML scalar mistake like allowed_openai_params: reasoning_effort would allowlist single characters.""" - assert litellm.Router._declared_param_allowlist( - {"allowed_openai_params": ["reasoning_effort", 3]} - ) == frozenset({"reasoning_effort"}) + assert litellm.Router._declared_param_allowlist({"allowed_openai_params": ["reasoning_effort", 3]}) == frozenset( + {"reasoning_effort"} + ) assert litellm.Router._declared_param_allowlist({"allowed_openai_params": "reasoning_effort"}) == frozenset() assert litellm.Router._declared_param_allowlist({}) == frozenset() @@ -12377,11 +12760,7 @@ class TestTierParamsTheTargetAccepts: @pytest.mark.parametrize( "deployment", - [ - {"model_name": "x"}, - {"model_name": "x", "litellm_params": {}}, - {"model_name": "x", "litellm_params": {"model": "not-a-real-provider/nope"}}, - ], + [{"model_name": "x"}, {"model_name": "x", "litellm_params": {}}, {"model_name": "x", "litellm_params": {"model": "not-a-real-provider/nope"}}], ) def test_deployment_accepts_param_fails_open(self, deployment): """An unresolvable deployment must not be the reason a param is dropped.""" @@ -12610,7 +12989,9 @@ class TestPreRoutingTierDrivesFallbacks: async def test_the_selected_tier_fallback_chain_runs(self): router = self._router([{"tier1": ["backup-a"]}]) - response = await router.acompletion(model="smart-router", messages=[{"role": "user", "content": "hi"}]) + response = await router.acompletion( + model="smart-router", messages=[{"role": "user", "content": "hi"}] + ) assert response.choices[0].message.content == "from backup-a" @@ -12643,7 +13024,9 @@ class TestPreRoutingTierDrivesFallbacks: async def test_a_request_without_a_pre_routing_hook_still_uses_its_own_group(self): router = self._router([{"tier1": ["backup-a"]}]) - response = await router.acompletion(model="tier1", messages=[{"role": "user", "content": "hi"}]) + response = await router.acompletion( + model="tier1", messages=[{"role": "user", "content": "hi"}] + ) assert response.choices[0].message.content == "from backup-a" From ff942c3a74e3d5a40e44f97111fffae56f29c06e Mon Sep 17 00:00:00 2001 From: moe-berri Date: Sat, 5 Sep 2026 11:05:50 -0700 Subject: [PATCH 11/13] fix(auto-router compression): surface a stored model-only policy in the edit form The backend treats either compression key on its own as an authoritative policy, but hydrate returned the untouched inherit state whenever the routing key was absent. A config carrying only auto_router_model_compression was therefore invisible in the form, and picking a routing value then overwrote the stored model hop. Only neither key set now reads as untouched, and an absent key on either hop hydrates as no compression for that hop rather than same-as-the-other. --- .../buildAutoRouterCompression.test.ts | 15 ++++++++++++++ .../add_model/buildAutoRouterCompression.ts | 20 ++++++++++++------- 2 files changed, 28 insertions(+), 7 deletions(-) diff --git a/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.test.ts b/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.test.ts index ea6eaf99106..20d6af50d18 100644 --- a/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.test.ts +++ b/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.test.ts @@ -96,6 +96,21 @@ describe("hydrateAutoRouterCompression", () => { 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)); diff --git a/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.ts b/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.ts index 47d0e3db7e4..6f401a12865 100644 --- a/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.ts +++ b/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.ts @@ -50,14 +50,20 @@ export const hydrateAutoRouterCompression = (litellmParams: { auto_router_routing_compression?: string | null; auto_router_model_compression?: string | null; }): AutoRouterCompressionState => { - const routing = litellmParams.auto_router_routing_compression ?? undefined; - if (routing === undefined) return DEFAULT_AUTO_ROUTER_COMPRESSION; + const storedRouting = litellmParams.auto_router_routing_compression ?? undefined; + const storedModel = litellmParams.auto_router_model_compression ?? undefined; - // An absent model key is no model-hop compression, not same-as-routing: the backend - // reads it as None (policy_from_litellm_params). Hydrating it as same-as-routing - // would make re-saving an unrelated edit write the routing guardrail onto the model - // hop and silently start compressing the model call. - const model = litellmParams.auto_router_model_compression ?? NO_COMPRESSION; + // 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 }; }; From 723bc2140fb55dc7f7fdf56cac184763e092c8ae Mon Sep 17 00:00:00 2001 From: moe-berri Date: Sat, 5 Sep 2026 11:19:39 -0700 Subject: [PATCH 12/13] refactor(auto-router compression): resolve the policy without a loop-local rebind The marker walk rebound a loop-local on each iteration, which is the mutation the repository's convention exists to discourage, but a `: Final` cannot express that inside a loop body: basedpyright rejects it outright with 'A Final variable cannot be assigned within a loop'. A lazy generator binds the name once per item and never rebinds it, so the first marker carrying a policy still wins and the rest are never read. --- litellm/proxy/guardrails/auto_router_compression.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/guardrails/auto_router_compression.py b/litellm/proxy/guardrails/auto_router_compression.py index b6f32c1c46d..ab3ba4011c7 100644 --- a/litellm/proxy/guardrails/auto_router_compression.py +++ b/litellm/proxy/guardrails/auto_router_compression.py @@ -117,11 +117,11 @@ def policy_for_model( # request does not carry describes a different slice of traffic, so falling back # to it would apply, say, an "eu" policy to a "us" request purely on config order. untagged: Final = tuple(params for params in markers if not params.get("tags")) - for params in (*tag_matched, *untagged): - policy = policy_from_litellm_params(params) - if policy is not None: - return policy - return None + # Lazily, so the first marker carrying a policy still wins and the rest are never + # read. A generator rather than a loop-local: the name is bound once per item and + # never rebound, which `: Final` cannot express inside a loop body. + 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: From a0b07b47912caf10488f0e7926807cc83f5611d7 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Sat, 5 Sep 2026 11:33:21 -0700 Subject: [PATCH 13/13] docs(auto-router compression): cut the explanatory comments back The module, its routing hook and its tests carried long prose rationale where the repository allows only concise comments for genuinely complex logic. Trimmed to the non-obvious reasons and dropped the rest; no logic or test behaviour changes. --- .../guardrails/auto_router_compression.py | 89 ++++++------------- litellm/router.py | 27 ++---- .../test_auto_router_compression.py | 59 ++++-------- 3 files changed, 49 insertions(+), 126 deletions(-) diff --git a/litellm/proxy/guardrails/auto_router_compression.py b/litellm/proxy/guardrails/auto_router_compression.py index ab3ba4011c7..98707e7ddca 100644 --- a/litellm/proxy/guardrails/auto_router_compression.py +++ b/litellm/proxy/guardrails/auto_router_compression.py @@ -1,15 +1,10 @@ """ -Decouples prompt compression between an auto router's routing decision and the -model it routes to. An auto router marker deployment may set -``auto_router_routing_compression`` and/or ``auto_router_model_compression`` in its -``litellm_params`` to name the compression guardrail that hop should use, or -``"none"`` to run no compression on that hop. Neither key set means the request's -own compression guardrails (key/team/model-level, or an "Always on" guardrail) -apply to both hops unchanged, exactly as before this feature existed. +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"``. -Once either key is set, this auto router is authoritative: every other compression -guardrail is suppressed for that request, and only these two settings decide what -each hop sees. +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 @@ -29,12 +24,8 @@ if TYPE_CHECKING: COMPRESSION_GUARDRAIL_PROVIDERS: Final = frozenset({"headroom", "compresr"}) _NO_COMPRESSION: Final = "none" -# Compression guardrails this request's auto router has switched off. Deliberately a -# ContextVar rather than a metadata key: `refresh_proxy_server_request_body_snapshot` -# copies metadata into `proxy_server_request.body`, which deployments persist to spend -# logs. A suppression list that reaches a log the caller can read is a list the caller -# can replay, which would let any request switch off a PII or content-filter guardrail. -# Nothing here is caller-supplied, so there is no marker to forge in the first place. +# 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() ) @@ -45,9 +36,8 @@ def suppressed_compression_guardrails() -> frozenset[str]: return _suppressed_compression_guardrails.get() -# Whether `arm_pre_call` actually armed a model-side compression guardrail for this -# request. 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. +# 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 ) @@ -95,10 +85,8 @@ def policy_for_model( ) -> AutoRouterCompressionPolicy | None: """The compression policy of the auto router marker `model_alias` resolves to. - Both the proxy's pre-call arming and the router's routing hook resolve the policy - through here, with the same tag rule, so an alias carrying several tag-scoped - markers can never suppress one marker's guardrail and then route under another - marker's policy. + 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 @@ -113,13 +101,9 @@ def policy_for_model( tag_matched: Final = tuple( params for params in markers if (tags := params.get("tags")) and requested.issuperset(frozenset(tags)) ) - # Only untagged markers may serve as the fallback. A marker scoped to tags this - # request does not carry describes a different slice of traffic, so falling back - # to it would apply, say, an "eu" policy to a "us" request purely on config order. + # 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")) - # Lazily, so the first marker carrying a policy still wins and the rest are never - # read. A generator rather than a loop-local: the name is bound once per item and - # never rebound, which `: Final` cannot express inside a loop body. + # 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) @@ -145,12 +129,8 @@ def _compression_guardrail_classes() -> tuple[type, ...]: def is_compression_guardrail(guardrail: object) -> bool: """Whether `guardrail` is an instance of a compression guardrail provider. - Both hops are validated through here. The two policy fields are operator-supplied - names and nothing else constrains them, so without this a name that resolves to an - ordinary guardrail would be handed the conversation and invoked: the routing hop - calls `apply_guardrail` directly, which POSTs the content wherever that guardrail - sends it, and the model hop is added to `metadata["guardrails"]`, which runs it even - when it is not `default_on`. + 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) @@ -185,9 +165,6 @@ async def arm_pre_call( if not isinstance(model_alias, str) or not model_alias: return - # Read-only until a policy is confirmed: creating the metadata bucket for every - # request, including the vast majority with no auto-router compression policy, - # would be an unwanted side effect of merely checking for one. 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 ) @@ -209,8 +186,7 @@ async def arm_pre_call( ) ) - # Only a name that resolves to a real compression guardrail may be armed: this adds - # it to `metadata["guardrails"]`, which runs it even when it is not `default_on`. + # 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() ) @@ -226,8 +202,7 @@ async def arm_pre_call( requested: Final = metadata.get("guardrails") existing: Final = tuple(requested) if isinstance(requested, (list, tuple)) else () if policy.model not in existing: - # A list, not a tuple: litellm_pre_call_utils tests this key with - # isinstance(..., list) and extends it, and would drop a tuple on the floor. + # 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 @@ -240,25 +215,17 @@ def _as_routing_messages( async def messages_for_routing( policy: AutoRouterCompressionPolicy | None, - # list[dict], not Sequence[Mapping]: the async_pre_routing_hook protocol in - # litellm/types/router.py types `messages` as list[dict[str, Any]]. + # 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`. + """Messages to use for a routing decision, per `policy.routing`. None means the + caller should route on whatever it already has. - Returns None when the caller should route on whatever messages it already has. - - Always reads the live messages, never a pre-guardrail copy of them. The routing - hop compresses through a real guardrail, which POSTs the text to an external - compression service, so it must see what every other guardrail has already done - to the request. Routing on a snapshot taken before the pre-call hook would send - a masking guardrail's own input straight back out of the proxy. - - The consequence, when the model hop compressed and the two hops differ: the - messages in hand are that guardrail's output, and there is no un-compressed copy - left to route on. The routing decision reads the compressed text in that one - combination rather than leaking the original. + 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 @@ -277,9 +244,6 @@ async def messages_for_routing( ) return _as_routing_messages(messages) - # apply_guardrail below hands this guardrail the conversation and it POSTs the - # content to whatever service backs it, so the name has to be a compression - # guardrail rather than any guardrail the operator happened to name. if not is_compression_guardrail(guardrail): verbose_proxy_logger.warning( "AutoRouter compression: guardrail '%s' is not a compression guardrail; routing on uncompressed messages", @@ -291,9 +255,8 @@ async def messages_for_routing( "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") - # A throwaway request_data: apply_guardrail writes its stats onto this dict, not the - # real request's metadata, so routing-side compression never double-counts against - # extract_compression_saved_tokens's model-savings accounting. + # Throwaway: 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, diff --git a/litellm/router.py b/litellm/router.py index 81de4572af8..51e2dbe38e4 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -13047,25 +13047,17 @@ class Router: team_id_from_request, ) - # Resolved through the same tag-aware lookup the proxy's pre-call arming used, - # so an alias carrying several tag-scoped markers cannot suppress one marker's - # guardrail and then route under a different marker's policy. + # 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), ) - # When both hops share the same compression, the model-side guardrail already - # ran in the proxy's ordinary pre-call hook and compressed `messages` in place - # (arm_pre_call armed it whether or not it is `default_on`); reuse that result - # for routing too instead of paying for a second compression call against the - # same content. - # - # Only the proxy calls arm_pre_call, so that reuse is conditional on it having - # actually run: on the SDK path nothing arms the model hop and nothing has - # compressed anything, and taking the shortcut there would skip both hops and - # silently serve the request with no compression at all. + # 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() ) @@ -13082,12 +13074,9 @@ class Router: input=input, specific_deployment=specific_deployment, ) - # The strategy only echoes back whatever `messages` it was handed, so a - # routing-only compression must not leak into the response: the model call - # and downstream deployment-context filtering both key off this field. - # Compared by value, not identity: PreRoutingHookResponse is a pydantic model, - # and pydantic reconstructs a validated list field rather than keeping the - # exact object passed in, even when nothing about it changed. + # 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 diff --git a/tests/test_litellm/proxy/guardrails/test_auto_router_compression.py b/tests/test_litellm/proxy/guardrails/test_auto_router_compression.py index 676b3ba2967..db2f94306fb 100644 --- a/tests/test_litellm/proxy/guardrails/test_auto_router_compression.py +++ b/tests/test_litellm/proxy/guardrails/test_auto_router_compression.py @@ -1,20 +1,4 @@ -""" -Unit tests for litellm.proxy.guardrails.auto_router_compression. - -Covers: -- policy_from_litellm_params: absent keys mean no policy; the "none" sentinel - normalizes to explicit no-compression within an active policy; is_same -- policy_for_model: finds the auto-router marker deployment for an alias, picks the - tag-scoped marker the request's tags actually match, and never falls back to a - marker scoped to tags the request does not carry -- arm_pre_call: no-op without a policy; suppresses active compression guardrails - through request-scoped state rather than metadata, which reaches spend logs a - caller can read; arms the model-side guardrail even when it isn't default_on -- messages_for_routing: no-op without a policy; compresses the live messages every - earlier guardrail has already rewritten, never a pre-guardrail copy of them; - never writes stats onto the caller's own request_kwargs (regression for - double-counted compression savings) -""" +"""Unit tests for litellm.proxy.guardrails.auto_router_compression.""" import json from typing import Any @@ -137,8 +121,7 @@ class TestPolicyForModel: assert policy == AutoRouterCompressionPolicy(routing="headroom-a", model=None) def test_a_marker_scoped_to_other_tags_is_never_the_fallback(self): - """Regression: an "eu" marker describes a different slice of traffic, so a "us" - request must not fall back to its policy just because it is configured first.""" + """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"]), @@ -149,8 +132,7 @@ class TestPolicyForModel: assert policy == AutoRouterCompressionPolicy(routing="headroom-default", model=None) def test_no_untagged_fallback_means_no_policy(self): - """With only tag-scoped markers and none matching, there is no policy to apply: - inheriting an unrelated slice's compression is worse than inheriting nothing.""" + """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 @@ -268,9 +250,8 @@ class TestArmPreCall: @pytest.mark.asyncio async def test_suppression_state_never_enters_request_metadata(self): - """Regression (security): a suppression list written to metadata is copied into - proxy_server_request.body and persisted to spend logs, so a caller could read it - back and replay it to switch off a PII or content-filter guardrail.""" + """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 @@ -312,9 +293,8 @@ class TestArmPreCall: @pytest.mark.asyncio async def test_arm_pre_call_keeps_no_copy_of_the_prompt(self): - """Regression (security): arm_pre_call runs before the pre-call guardrails, so - any copy of the messages it retained would be the pre-masking text. Routing-side - compression POSTs its input to an external service, so that copy must not exist.""" + """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"}]} @@ -337,10 +317,8 @@ class TestMessagesForRouting: @pytest.mark.asyncio async def test_routing_none_never_reaches_for_a_pre_guardrail_copy(self): - """Routing asked for no compression while the model hop compressed, so the - messages in hand are that guardrail's output and no uncompressed copy survives. - Routing reads them as-is: the alternative is keeping a pre-guardrail copy, which - is the text a masking guardrail exists to remove.""" + """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"}] @@ -362,10 +340,8 @@ class TestMessagesForRouting: @pytest.mark.asyncio async def test_routing_compresses_what_the_other_guardrails_left_behind(self, registered_guardrail): - """Regression (security): routing-side compression POSTs its input to an external - service, so it must read the live messages every earlier guardrail has already - rewritten. Routing on a pre-guardrail copy would send a masking guardrail's own - input straight back out of the proxy.""" + """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]"}] @@ -376,10 +352,8 @@ class TestMessagesForRouting: @pytest.mark.asyncio async def test_a_non_compression_guardrail_is_never_invoked_for_routing(self, monkeypatch): - """Regression (security): the policy fields are operator-supplied names that - nothing else constrains. apply_guardrail hands the guardrail the conversation - and it POSTs that content to whatever service backs it, so naming an ordinary - guardrail must not turn the routing hop into a way to ship prompts there.""" + """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") @@ -397,11 +371,8 @@ class TestMessagesForRouting: @pytest.mark.asyncio async def test_guardrail_receives_a_throwaway_request_data_not_the_real_request_kwargs(self, registered_guardrail): - """Regression: a real compression guardrail writes its stats onto whatever - `request_data` dict it's given (`add_standard_logging_guardrail_information_to_ - request_data`). If that were the caller's own `request_kwargs`, routing-side - compression would double-count into extract_compression_saved_tokens, which - sums every guardrail_information entry on the real request's metadata.""" + """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": {}}