mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-08 22:21:35 +00:00
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.
This commit is contained in:
moe-berri
parent
dd60b7e40f
commit
9e286fe94b
8 changed files with 322 additions and 133 deletions
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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"}]
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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<CompressionControlsProps> = ({
|
|||
}) => {
|
||||
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];
|
||||
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue