scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-10 22:41:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 83

refactor(auto-router compression): satisfy the LIT001/LIT002 type-discipline gate

Browse source 
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.
This commit is contained in:
moe-berri 2026-09-04 17:42:52 -07:00
parent e273cf301f
commit dc63428395
5 changed files with 83 additions and 80 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

8
litellm/integrations/custom_guardrail.py
View file

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

2
litellm/proxy/common_request_processing.py
View file

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

85
litellm/proxy/guardrails/auto_router_compression.py
View file

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

3
litellm/router.py
View file

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

65
tests/test_litellm/proxy/guardrails/test_auto_router_compression.py
View file

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