mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-08 22:21:35 +00:00
Merge pull request #39823 from BerriAI/litellm_auto_router_compression_split
feat(auto-router): decouple compression between the routing decision and the model call
This commit is contained in:
moe-berri
GitHub
commit
b3f28a77d8
21 changed files with 1650 additions and 223 deletions
|
|
@ -948,6 +948,23 @@ class CustomGuardrail(CustomLogger):
|
|||
"""
|
||||
return False
|
||||
|
||||
def _suppressed_by_auto_router_compression(self) -> bool:
|
||||
"""True when an auto router's own compression policy suppresses this guardrail.
|
||||
|
||||
Reads request-scoped state set by `arm_pre_call`, never request metadata. The
|
||||
caller controls metadata, and metadata reaches spend logs the caller can read,
|
||||
so a suppression list carried there would be one a request could replay to
|
||||
switch off a PII or content-filter guardrail for itself.
|
||||
"""
|
||||
name: Final = self.guardrail_name
|
||||
if not name:
|
||||
return False
|
||||
from litellm.proxy.guardrails.auto_router_compression import (
|
||||
suppressed_compression_guardrails,
|
||||
)
|
||||
|
||||
return name in suppressed_compression_guardrails()
|
||||
|
||||
def should_run_guardrail(
|
||||
self,
|
||||
data,
|
||||
|
|
@ -956,6 +973,9 @@ class CustomGuardrail(CustomLogger):
|
|||
"""
|
||||
Returns True if the guardrail should be run on the event_type
|
||||
"""
|
||||
if self._suppressed_by_auto_router_compression():
|
||||
return False
|
||||
|
||||
requested_guardrails: Final = self.get_guardrail_from_metadata(data)
|
||||
disable_global_guardrail: Final = self.get_disable_global_guardrail(data)
|
||||
opted_out_global_guardrails: Final = self.get_opted_out_global_guardrails_from_metadata(data)
|
||||
|
|
|
|||
|
|
@ -61,6 +61,7 @@ from litellm.proxy.common_utils.sse_keepalive import (
|
|||
wrap_sse_stream_with_keepalive_pings,
|
||||
)
|
||||
from litellm.proxy.dd_span_tagger import DDSpanTagger
|
||||
from litellm.proxy.guardrails.auto_router_compression import arm_pre_call as _arm_auto_router_compression
|
||||
from litellm.proxy.route_llm_request import route_request
|
||||
from litellm.proxy.utils import ProxyLogging, _check_and_merge_model_level_guardrails
|
||||
from litellm.router import Router
|
||||
|
|
@ -2004,6 +2005,12 @@ class ProxyBaseLLMRequestProcessing:
|
|||
trust_client_model_info=False,
|
||||
)
|
||||
|
||||
# An auto router with its own compression policy is authoritative for this
|
||||
# request: suppress every other compression guardrail and arm whichever one
|
||||
# the policy names for the model call, before those guardrails get a chance
|
||||
# to run below.
|
||||
await _arm_auto_router_compression(data=self.data, llm_router=llm_router)
|
||||
|
||||
self.data = await proxy_logging_obj.pre_call_hook(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
data=self.data,
|
||||
|
|
|
|||
267
litellm/proxy/guardrails/auto_router_compression.py
Normal file
267
litellm/proxy/guardrails/auto_router_compression.py
Normal file
|
|
@ -0,0 +1,267 @@
|
|||
"""
|
||||
Decouples prompt compression between an auto router's routing decision and the model
|
||||
it routes to, via ``auto_router_routing_compression`` / ``auto_router_model_compression``
|
||||
on the marker deployment: a guardrail name, or ``"none"``.
|
||||
|
||||
Neither key set inherits today's behaviour. Either key set makes the auto router
|
||||
authoritative and suppresses every other compression guardrail for that request.
|
||||
"""
|
||||
|
||||
import contextvars
|
||||
from collections.abc import Iterable, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.litellm_core_utils.core_helpers import get_or_create_metadata_bucket
|
||||
from litellm.router_utils.auto_router_model_naming import AUTO_ROUTER_MODEL_PREFIX
|
||||
from litellm.types.utils import GenericGuardrailAPIInputs
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.integrations.custom_guardrail import CustomGuardrail
|
||||
from litellm.router import Router
|
||||
|
||||
COMPRESSION_GUARDRAIL_PROVIDERS: Final = frozenset({"headroom", "compresr"})
|
||||
_NO_COMPRESSION: Final = "none"
|
||||
|
||||
# A ContextVar, not metadata: metadata reaches spend logs the caller can read, and a
|
||||
# suppression list they can read is one they can replay to disable any guardrail.
|
||||
_suppressed_compression_guardrails: Final[contextvars.ContextVar[frozenset[str]]] = contextvars.ContextVar(
|
||||
"litellm_auto_router_suppressed_compression_guardrails", default=frozenset()
|
||||
)
|
||||
|
||||
|
||||
def suppressed_compression_guardrails() -> frozenset[str]:
|
||||
"""Names of the compression guardrails this request's auto router suppresses."""
|
||||
return _suppressed_compression_guardrails.get()
|
||||
|
||||
|
||||
# Only the proxy calls `arm_pre_call`, so on the SDK path nothing arms and nothing
|
||||
# compresses; the router must not assume the model hop already ran.
|
||||
_model_hop_armed: Final[contextvars.ContextVar[bool]] = contextvars.ContextVar(
|
||||
"litellm_auto_router_model_hop_armed", default=False
|
||||
)
|
||||
|
||||
|
||||
def model_hop_compression_armed() -> bool:
|
||||
"""True when this request's model-side compression guardrail was actually armed."""
|
||||
return _model_hop_armed.get()
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class AutoRouterCompressionPolicy:
|
||||
"""An auto router's compression choice for each hop. ``None`` means no compression."""
|
||||
|
||||
routing: str | None
|
||||
model: str | None
|
||||
|
||||
@property
|
||||
def is_same(self) -> bool:
|
||||
return self.routing == self.model
|
||||
|
||||
|
||||
def _normalized_compression_choice(raw: object) -> str | None:
|
||||
if not isinstance(raw, str) or not raw:
|
||||
return None
|
||||
return None if raw.strip().lower() == _NO_COMPRESSION else raw
|
||||
|
||||
|
||||
def policy_from_litellm_params(litellm_params: Mapping[str, object]) -> AutoRouterCompressionPolicy | None:
|
||||
raw_routing: Final = litellm_params.get("auto_router_routing_compression")
|
||||
raw_model: Final = litellm_params.get("auto_router_model_compression")
|
||||
if raw_routing is None and raw_model is None:
|
||||
return None
|
||||
return AutoRouterCompressionPolicy(
|
||||
routing=_normalized_compression_choice(raw_routing),
|
||||
model=_normalized_compression_choice(raw_model),
|
||||
)
|
||||
|
||||
|
||||
def policy_for_model(
|
||||
llm_router: "Router | None",
|
||||
model_alias: str,
|
||||
team_id: str | None,
|
||||
request_tags: Sequence[str],
|
||||
) -> AutoRouterCompressionPolicy | None:
|
||||
"""The compression policy of the auto router marker `model_alias` resolves to.
|
||||
|
||||
Pre-call arming and the routing hook both resolve through here, so an alias with
|
||||
several tag-scoped markers cannot suppress under one and then route under another.
|
||||
"""
|
||||
if llm_router is None:
|
||||
return None
|
||||
deployments: Final = llm_router.get_model_list(model_name=model_alias, team_id=team_id) or ()
|
||||
markers: Final = tuple(
|
||||
litellm_params
|
||||
for deployment in deployments
|
||||
if isinstance(litellm_params := deployment.get("litellm_params"), Mapping) # pyright: ignore[reportUnnecessaryIsInstance] # filters out non-Mapping
|
||||
and str(litellm_params.get("model", "")).startswith(AUTO_ROUTER_MODEL_PREFIX)
|
||||
)
|
||||
requested: Final = frozenset(request_tags)
|
||||
tag_matched: Final = tuple(
|
||||
params for params in markers if (tags := params.get("tags")) and requested.issuperset(frozenset(tags))
|
||||
)
|
||||
# Untagged only: a marker scoped to tags this request lacks describes other traffic.
|
||||
untagged: Final = tuple(params for params in markers if not params.get("tags"))
|
||||
# Lazy, so the first marker carrying a policy wins and the rest are never read.
|
||||
candidates: Final = (policy_from_litellm_params(params) for params in (*tag_matched, *untagged))
|
||||
return next((policy for policy in candidates if policy is not None), None)
|
||||
|
||||
|
||||
def team_id_from_request(request_kwargs: Mapping[str, object]) -> str | None:
|
||||
"""The caller's team id, from whichever metadata bucket this surface writes to."""
|
||||
for meta_key in ("metadata", "litellm_metadata"):
|
||||
meta = request_kwargs.get(meta_key)
|
||||
if isinstance(meta, Mapping):
|
||||
team_id = meta.get("user_api_key_team_id")
|
||||
if isinstance(team_id, str):
|
||||
return team_id
|
||||
return None
|
||||
|
||||
|
||||
def _compression_guardrail_classes() -> tuple[type, ...]:
|
||||
"""The registered guardrail classes whose provider compresses prompts."""
|
||||
from litellm.proxy.guardrails.guardrail_registry import guardrail_class_registry
|
||||
|
||||
return tuple(cls for name, cls in guardrail_class_registry.items() if name in COMPRESSION_GUARDRAIL_PROVIDERS)
|
||||
|
||||
|
||||
def is_compression_guardrail(guardrail: object) -> bool:
|
||||
"""Whether `guardrail` is an instance of a compression guardrail provider.
|
||||
|
||||
Both hops validate through here: the policy fields are operator-supplied names, and
|
||||
an unvalidated one would get handed the conversation and invoked.
|
||||
"""
|
||||
classes: Final = _compression_guardrail_classes()
|
||||
return bool(classes) and isinstance(guardrail, classes)
|
||||
|
||||
|
||||
def _active_compression_guardrails() -> tuple["CustomGuardrail", ...]:
|
||||
"""Every currently-active guardrail whose type is a compression guardrail."""
|
||||
import litellm
|
||||
from litellm.integrations.custom_guardrail import CustomGuardrail
|
||||
|
||||
if not _compression_guardrail_classes():
|
||||
return ()
|
||||
active: Final = litellm.logging_callback_manager.get_custom_loggers_for_type(callback_type=CustomGuardrail)
|
||||
return tuple(cb for cb in active if is_compression_guardrail(cb) and cb.guardrail_name)
|
||||
|
||||
|
||||
async def arm_pre_call(
|
||||
data: dict[str, object], # mutable-ok: arms the live request dict in place
|
||||
llm_router: "Router | None",
|
||||
) -> None:
|
||||
"""Apply an auto router's compression policy, if any, before guardrails run.
|
||||
|
||||
Suppresses every other compression guardrail and re-enables the model-side
|
||||
guardrail the policy names (if any) even when it isn't ``default_on``.
|
||||
"""
|
||||
_suppressed_compression_guardrails.set(frozenset())
|
||||
_model_hop_armed.set(False)
|
||||
if llm_router is None:
|
||||
return
|
||||
|
||||
model_alias: Final = data.get("model")
|
||||
if not isinstance(model_alias, str) or not model_alias:
|
||||
return
|
||||
|
||||
from litellm.router_strategy.tag_based_routing import (
|
||||
_get_tags_from_request_kwargs, # pyright: ignore[reportPrivateUsage] # used in router.py and budget_limiter.py too
|
||||
)
|
||||
|
||||
policy: Final = policy_for_model(
|
||||
llm_router=llm_router,
|
||||
model_alias=model_alias,
|
||||
team_id=team_id_from_request(data),
|
||||
request_tags=_get_tags_from_request_kwargs(data),
|
||||
)
|
||||
if policy is None:
|
||||
return
|
||||
|
||||
_suppressed_compression_guardrails.set(
|
||||
frozenset(
|
||||
name
|
||||
for guardrail in _active_compression_guardrails()
|
||||
if (name := guardrail.guardrail_name) and name != policy.model
|
||||
)
|
||||
)
|
||||
|
||||
# Arming adds the name to `metadata["guardrails"]`, which runs it even if not default_on.
|
||||
armed_model_hop: Final = policy.model is not None and any(
|
||||
guardrail.guardrail_name == policy.model for guardrail in _active_compression_guardrails()
|
||||
)
|
||||
if policy.model is not None and not armed_model_hop:
|
||||
verbose_proxy_logger.warning(
|
||||
"AutoRouter compression: '%s' is not an active compression guardrail; the model hop is uncompressed",
|
||||
policy.model,
|
||||
)
|
||||
|
||||
if armed_model_hop:
|
||||
_model_hop_armed.set(True)
|
||||
_, metadata = get_or_create_metadata_bucket(data)
|
||||
requested: Final = metadata.get("guardrails")
|
||||
existing: Final = tuple(requested) if isinstance(requested, (list, tuple)) else ()
|
||||
if policy.model not in existing:
|
||||
# A list: litellm_pre_call_utils isinstance-checks this key and drops a tuple.
|
||||
metadata["guardrails"] = [*existing, policy.model] # mutable-ok: this key's contract is a list
|
||||
|
||||
|
||||
def _as_routing_messages(
|
||||
messages: Iterable[Mapping[str, object]],
|
||||
) -> list[dict[str, object]]: # mutable-ok: shape fixed by the pre-routing hook protocol
|
||||
"""A fresh, independently mutable copy, the shape the pre-routing hook takes."""
|
||||
return [dict(message) for message in messages] # mutable-ok: shape fixed by the pre-routing hook protocol
|
||||
|
||||
|
||||
async def messages_for_routing(
|
||||
policy: AutoRouterCompressionPolicy | None,
|
||||
# list[dict], not Sequence[Mapping]: fixed by the async_pre_routing_hook protocol.
|
||||
messages: list[dict[str, object]] | None, # mutable-ok: shape fixed by the pre-routing hook protocol
|
||||
request_kwargs: Mapping[str, object],
|
||||
) -> list[dict[str, object]] | None: # mutable-ok: shape fixed by the pre-routing hook protocol
|
||||
"""Messages to use for a routing decision, per `policy.routing`. None means the
|
||||
caller should route on whatever it already has.
|
||||
|
||||
Reads the live messages, never a pre-guardrail copy: this compresses through a real
|
||||
guardrail that POSTs the text out, so routing on a pre-masking snapshot would leak
|
||||
what the masking guardrail stripped. When the model hop already compressed and the
|
||||
hops differ, routing therefore reads the compressed text rather than the original.
|
||||
"""
|
||||
if policy is None or policy.routing is None:
|
||||
return None
|
||||
|
||||
if not messages:
|
||||
return None
|
||||
|
||||
from litellm.proxy.common_utils.registry_read_through import (
|
||||
get_initialized_guardrail_with_read_through,
|
||||
)
|
||||
|
||||
guardrail: Final = await get_initialized_guardrail_with_read_through(policy.routing)
|
||||
if guardrail is None:
|
||||
verbose_proxy_logger.warning(
|
||||
"AutoRouter compression: guardrail '%s' not found; routing on uncompressed messages", policy.routing
|
||||
)
|
||||
return _as_routing_messages(messages)
|
||||
|
||||
if not is_compression_guardrail(guardrail):
|
||||
verbose_proxy_logger.warning(
|
||||
"AutoRouter compression: guardrail '%s' is not a compression guardrail; routing on uncompressed messages",
|
||||
policy.routing,
|
||||
)
|
||||
return _as_routing_messages(messages)
|
||||
|
||||
inputs: Final[GenericGuardrailAPIInputs] = {
|
||||
"structured_messages": _as_routing_messages(messages) # pyright: ignore[reportAssignmentType] # plain dicts, not AllMessageValues; see headroom.py's own use of this shape
|
||||
}
|
||||
model: Final = request_kwargs.get("model")
|
||||
# Throwaway: apply_guardrail writes stats here, so routing never double-counts into
|
||||
# extract_compression_saved_tokens.
|
||||
stats_sink: Final = {"messages": messages, "model": model} # mutable-ok: apply_guardrail writes its stats here
|
||||
result: Final = await guardrail.apply_guardrail(
|
||||
inputs=inputs,
|
||||
request_data=stats_sink,
|
||||
input_type="request",
|
||||
)
|
||||
compressed: Final = result.get("structured_messages")
|
||||
return compressed if isinstance(compressed, list) else _as_routing_messages(messages)
|
||||
|
|
@ -8636,7 +8636,7 @@ class Router:
|
|||
raise ValueError(ptu_error)
|
||||
zeroed_pricing: Final = zeroed_ptu_pricing(_model_info, _litellm_params) if config_sourced else None
|
||||
litellm_params: Final[LiteLLM_Params] = LiteLLM_Params(
|
||||
**(
|
||||
**( # pyright: ignore[reportArgumentType] # untyped merged dict; already true for every field here
|
||||
_litellm_params
|
||||
if zeroed_pricing is None
|
||||
else MappingProxyType({**_litellm_params, **zeroed_pricing})
|
||||
|
|
@ -13033,13 +13033,48 @@ class Router:
|
|||
)
|
||||
return None
|
||||
|
||||
pre_routing_hook_response: Final = await selected_strategy.strategy.async_pre_routing_hook(
|
||||
from litellm.proxy.guardrails.auto_router_compression import (
|
||||
messages_for_routing,
|
||||
model_hop_compression_armed,
|
||||
policy_for_model,
|
||||
team_id_from_request,
|
||||
)
|
||||
|
||||
# Same tag-aware lookup the proxy's pre-call arming used, so an alias with
|
||||
# several tag-scoped markers cannot suppress under one and route under another.
|
||||
compression_policy: Final = policy_for_model(
|
||||
llm_router=self,
|
||||
model_alias=registered_model_name,
|
||||
team_id=team_id_from_request(request_kwargs),
|
||||
request_tags=_get_tags_from_request_kwargs(request_kwargs),
|
||||
)
|
||||
# Shared compression already ran in the pre-call hook, so reuse it rather than
|
||||
# compressing twice. Conditional on arming having actually happened: only the
|
||||
# proxy arms, and on the SDK path the shortcut would skip both hops entirely.
|
||||
needs_independent_routing_compression: Final = compression_policy is not None and not (
|
||||
compression_policy.is_same and compression_policy.model is not None and model_hop_compression_armed()
|
||||
)
|
||||
routing_messages: Final = (
|
||||
await messages_for_routing(policy=compression_policy, messages=messages, request_kwargs=request_kwargs)
|
||||
if needs_independent_routing_compression
|
||||
else None
|
||||
)
|
||||
|
||||
routed: Final = await selected_strategy.strategy.async_pre_routing_hook(
|
||||
model=registered_model_name,
|
||||
request_kwargs=request_kwargs,
|
||||
messages=messages,
|
||||
messages=routing_messages if routing_messages is not None else messages,
|
||||
input=input,
|
||||
specific_deployment=specific_deployment,
|
||||
)
|
||||
# Routing-only compression must not leak into the response: the model call and
|
||||
# deployment-context filtering key off this field. Compared by value, since
|
||||
# pydantic rebuilds the list rather than keeping the object passed in.
|
||||
pre_routing_hook_response: Final = (
|
||||
routed.model_copy(update={"messages": messages}) # mutable-ok: pydantic's model_copy takes a dict
|
||||
if routed is not None and routing_messages is not None and routed.messages == routing_messages
|
||||
else routed
|
||||
)
|
||||
self._record_routing_decision(
|
||||
request_kwargs=request_kwargs,
|
||||
routing_decision=(pre_routing_hook_response.routing_decision if pre_routing_hook_response else None),
|
||||
|
|
|
|||
|
|
@ -361,6 +361,10 @@ class GenericLiteLLMParams(CredentialLiteLLMParams, CustomPricingLiteLLMParams):
|
|||
auto_router_default_model: str | None = None
|
||||
auto_router_embedding_model: str | None = None
|
||||
auto_router_max_input_chars: int | None = None
|
||||
# Compression policy for the two hops of a routed request. Both unset means the
|
||||
# request's own compression guardrails apply to both, as they always have.
|
||||
auto_router_routing_compression: str | None = None
|
||||
auto_router_model_compression: str | None = None
|
||||
|
||||
# complexity-router params
|
||||
complexity_router_config: dict | None = None
|
||||
|
|
|
|||
|
|
@ -3769,6 +3769,8 @@ all_litellm_params = (
|
|||
"auto_router_default_model",
|
||||
"auto_router_embedding_model",
|
||||
"auto_router_max_input_chars",
|
||||
"auto_router_routing_compression",
|
||||
"auto_router_model_compression",
|
||||
"complexity_router_config",
|
||||
"complexity_router_default_model",
|
||||
"adaptive_router_config",
|
||||
|
|
|
|||
|
|
@ -17,7 +17,6 @@ if TYPE_CHECKING:
|
|||
|
||||
|
||||
class TestCustomGuardrailDeploymentHook:
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_pre_call_deployment_hook_no_guardrails(self):
|
||||
"""Test that method returns kwargs unchanged when no guardrails are present"""
|
||||
|
|
@ -30,18 +29,14 @@ class TestCustomGuardrailDeploymentHook:
|
|||
"guardrails": None,
|
||||
}
|
||||
|
||||
result = await custom_guardrail.async_pre_call_deployment_hook(
|
||||
kwargs=kwargs, call_type=CallTypes.completion
|
||||
)
|
||||
result = await custom_guardrail.async_pre_call_deployment_hook(kwargs=kwargs, call_type=CallTypes.completion)
|
||||
|
||||
assert result == kwargs
|
||||
|
||||
# Test with guardrails as non-list
|
||||
kwargs["guardrails"] = "not_a_list"
|
||||
|
||||
result = await custom_guardrail.async_pre_call_deployment_hook(
|
||||
kwargs=kwargs, call_type=CallTypes.completion
|
||||
)
|
||||
result = await custom_guardrail.async_pre_call_deployment_hook(kwargs=kwargs, call_type=CallTypes.completion)
|
||||
|
||||
assert result == kwargs
|
||||
|
||||
|
|
@ -68,9 +63,7 @@ class TestCustomGuardrailDeploymentHook:
|
|||
"user_api_key_request_route": "test_route",
|
||||
}
|
||||
|
||||
result = await custom_guardrail.async_pre_call_deployment_hook(
|
||||
kwargs=kwargs, call_type=CallTypes.completion
|
||||
)
|
||||
result = await custom_guardrail.async_pre_call_deployment_hook(kwargs=kwargs, call_type=CallTypes.completion)
|
||||
|
||||
# Verify async_pre_call_hook was called with correct parameters
|
||||
custom_guardrail.async_pre_call_hook.assert_called_once()
|
||||
|
|
@ -103,9 +96,7 @@ class TestCustomGuardrailDeploymentHook:
|
|||
super().__init__(guardrail_name="g1", default_on=True)
|
||||
self.pre_call_count = 0
|
||||
|
||||
async def async_pre_call_hook(
|
||||
self, user_api_key_dict, cache, data, call_type
|
||||
):
|
||||
async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type):
|
||||
self.pre_call_count += 1
|
||||
return data
|
||||
|
||||
|
|
@ -118,9 +109,7 @@ class TestCustomGuardrailDeploymentHook:
|
|||
}
|
||||
|
||||
guardrail.mark_pre_call_hook_ran(kwargs)
|
||||
await guardrail.async_pre_call_deployment_hook(
|
||||
kwargs=kwargs, call_type=CallTypes.completion
|
||||
)
|
||||
await guardrail.async_pre_call_deployment_hook(kwargs=kwargs, call_type=CallTypes.completion)
|
||||
|
||||
assert guardrail.pre_call_count == 0
|
||||
|
||||
|
|
@ -134,9 +123,7 @@ class TestCustomGuardrailDeploymentHook:
|
|||
super().__init__(guardrail_name="g1", default_on=True)
|
||||
self.pre_call_count = 0
|
||||
|
||||
async def async_pre_call_hook(
|
||||
self, user_api_key_dict, cache, data, call_type
|
||||
):
|
||||
async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type):
|
||||
self.pre_call_count += 1
|
||||
return data
|
||||
|
||||
|
|
@ -148,9 +135,7 @@ class TestCustomGuardrailDeploymentHook:
|
|||
"metadata": {},
|
||||
}
|
||||
|
||||
await guardrail.async_pre_call_deployment_hook(
|
||||
kwargs=kwargs, call_type=CallTypes.completion
|
||||
)
|
||||
await guardrail.async_pre_call_deployment_hook(kwargs=kwargs, call_type=CallTypes.completion)
|
||||
|
||||
assert guardrail.pre_call_count == 1
|
||||
|
||||
|
|
@ -179,9 +164,7 @@ class TestCustomGuardrailDeploymentHook:
|
|||
super().__init__(guardrail_name="g1", default_on=True)
|
||||
self.pre_call_count = 0
|
||||
|
||||
async def async_pre_call_hook(
|
||||
self, user_api_key_dict, cache, data, call_type
|
||||
):
|
||||
async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type):
|
||||
self.pre_call_count += 1
|
||||
return data
|
||||
|
||||
|
|
@ -193,15 +176,12 @@ class TestCustomGuardrailDeploymentHook:
|
|||
"metadata": {PRE_CALL_EXECUTED_GUARDRAILS_KEY: ["g1"]},
|
||||
}
|
||||
|
||||
await guardrail.async_pre_call_deployment_hook(
|
||||
kwargs=kwargs, call_type=CallTypes.completion
|
||||
)
|
||||
await guardrail.async_pre_call_deployment_hook(kwargs=kwargs, call_type=CallTypes.completion)
|
||||
|
||||
assert guardrail.pre_call_count == 1
|
||||
|
||||
|
||||
class TestCustomGuardrailShouldRunGuardrail:
|
||||
|
||||
def test_should_run_guardrail_with_litellm_metadata(self):
|
||||
"""Test that should_run_guardrail works with litellm_metadata pattern"""
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
|
|
@ -218,9 +198,7 @@ class TestCustomGuardrailShouldRunGuardrail:
|
|||
"litellm_metadata": {"guardrails": ["test_guardrail"]},
|
||||
}
|
||||
|
||||
result = custom_guardrail.should_run_guardrail(
|
||||
data=data, event_type=GuardrailEventHooks.pre_call
|
||||
)
|
||||
result = custom_guardrail.should_run_guardrail(data=data, event_type=GuardrailEventHooks.pre_call)
|
||||
|
||||
assert result is True
|
||||
|
||||
|
|
@ -240,9 +218,7 @@ class TestCustomGuardrailShouldRunGuardrail:
|
|||
"metadata": {"guardrails": ["test_guardrail"]},
|
||||
}
|
||||
|
||||
result = custom_guardrail.should_run_guardrail(
|
||||
data=data, event_type=GuardrailEventHooks.pre_call
|
||||
)
|
||||
result = custom_guardrail.should_run_guardrail(data=data, event_type=GuardrailEventHooks.pre_call)
|
||||
|
||||
assert result is True
|
||||
|
||||
|
|
@ -259,9 +235,7 @@ class TestCustomGuardrailShouldRunGuardrail:
|
|||
# Test with guardrails at root level
|
||||
data = {"model": "gpt-3.5-turbo", "guardrails": ["test_guardrail"]}
|
||||
|
||||
result = custom_guardrail.should_run_guardrail(
|
||||
data=data, event_type=GuardrailEventHooks.pre_call
|
||||
)
|
||||
result = custom_guardrail.should_run_guardrail(data=data, event_type=GuardrailEventHooks.pre_call)
|
||||
|
||||
assert result is True
|
||||
|
||||
|
|
@ -281,9 +255,7 @@ class TestCustomGuardrailShouldRunGuardrail:
|
|||
"litellm_metadata": {"guardrails": ["different_guardrail"]},
|
||||
}
|
||||
|
||||
result = custom_guardrail.should_run_guardrail(
|
||||
data=data, event_type=GuardrailEventHooks.pre_call
|
||||
)
|
||||
result = custom_guardrail.should_run_guardrail(data=data, event_type=GuardrailEventHooks.pre_call)
|
||||
|
||||
assert result is False
|
||||
|
||||
|
|
@ -302,9 +274,7 @@ class TestCustomGuardrailShouldRunGuardrail:
|
|||
"model": "gpt-3.5-turbo",
|
||||
"messages": [{"role": "user", "content": "test"}],
|
||||
}
|
||||
result = custom_guardrail.should_run_guardrail(
|
||||
data=data, event_type=GuardrailEventHooks.pre_call
|
||||
)
|
||||
result = custom_guardrail.should_run_guardrail(data=data, event_type=GuardrailEventHooks.pre_call)
|
||||
assert result is True, "Global guardrail should run when default_on=True"
|
||||
|
||||
# Test 2: User-injected disable at root level is IGNORED
|
||||
|
|
@ -316,9 +286,7 @@ class TestCustomGuardrailShouldRunGuardrail:
|
|||
result = custom_guardrail.should_run_guardrail(
|
||||
data=data_with_disable_root, event_type=GuardrailEventHooks.pre_call
|
||||
)
|
||||
assert (
|
||||
result is True
|
||||
), "User-injected disable_global_guardrails should be ignored"
|
||||
assert result is True, "User-injected disable_global_guardrails should be ignored"
|
||||
|
||||
# Test 3: User-injected disable in metadata is IGNORED
|
||||
data_with_disable_metadata = {
|
||||
|
|
@ -349,12 +317,8 @@ class TestCustomGuardrailShouldRunGuardrail:
|
|||
"metadata": {"user_api_key_metadata": {"disable_global_guardrails": True}},
|
||||
"litellm_metadata": {"request_tags": ["user-supplied"]},
|
||||
}
|
||||
result = custom_guardrail.should_run_guardrail(
|
||||
data=data_cross_key, event_type=GuardrailEventHooks.pre_call
|
||||
)
|
||||
assert (
|
||||
result is False
|
||||
), "Admin config in metadata must not be shadowed by user-supplied litellm_metadata"
|
||||
result = custom_guardrail.should_run_guardrail(data=data_cross_key, event_type=GuardrailEventHooks.pre_call)
|
||||
assert result is False, "Admin config in metadata must not be shadowed by user-supplied litellm_metadata"
|
||||
|
||||
# Test 6: After the pre-call strip runs, user-injected
|
||||
# user_api_key_metadata in the non-authoritative metadata key is gone.
|
||||
|
|
@ -365,12 +329,8 @@ class TestCustomGuardrailShouldRunGuardrail:
|
|||
"metadata": {"user_api_key_metadata": {"disable_global_guardrails": True}},
|
||||
"litellm_metadata": {}, # post-strip: attacker payload removed
|
||||
}
|
||||
result = custom_guardrail.should_run_guardrail(
|
||||
data=data_post_strip, event_type=GuardrailEventHooks.pre_call
|
||||
)
|
||||
assert (
|
||||
result is False
|
||||
), "Admin config in metadata must be respected when other metadata key is empty"
|
||||
result = custom_guardrail.should_run_guardrail(data=data_post_strip, event_type=GuardrailEventHooks.pre_call)
|
||||
assert result is False, "Admin config in metadata must be respected when other metadata key is empty"
|
||||
|
||||
def test_should_run_guardrail_key_disable_global_not_overruled_by_team_guardrail_list(
|
||||
self,
|
||||
|
|
@ -436,12 +396,7 @@ class TestCustomGuardrailShouldRunGuardrail:
|
|||
"messages": [{"role": "user", "content": "test"}],
|
||||
"opted_out_global_guardrails": ["global_guardrail"],
|
||||
}
|
||||
assert (
|
||||
custom_guardrail.should_run_guardrail(
|
||||
data=data_root, event_type=GuardrailEventHooks.pre_call
|
||||
)
|
||||
is True
|
||||
)
|
||||
assert custom_guardrail.should_run_guardrail(data=data_root, event_type=GuardrailEventHooks.pre_call) is True
|
||||
|
||||
# Test 2: User-injected opt-out in metadata is IGNORED
|
||||
data_metadata = {
|
||||
|
|
@ -450,10 +405,7 @@ class TestCustomGuardrailShouldRunGuardrail:
|
|||
"metadata": {"opted_out_global_guardrails": ["global_guardrail"]},
|
||||
}
|
||||
assert (
|
||||
custom_guardrail.should_run_guardrail(
|
||||
data=data_metadata, event_type=GuardrailEventHooks.pre_call
|
||||
)
|
||||
is True
|
||||
custom_guardrail.should_run_guardrail(data=data_metadata, event_type=GuardrailEventHooks.pre_call) is True
|
||||
)
|
||||
|
||||
# Test 4: a different guardrail in the opt-out list → still runs
|
||||
|
|
@ -462,12 +414,7 @@ class TestCustomGuardrailShouldRunGuardrail:
|
|||
"messages": [{"role": "user", "content": "test"}],
|
||||
"metadata": {"opted_out_global_guardrails": ["some_other_guardrail"]},
|
||||
}
|
||||
assert (
|
||||
custom_guardrail.should_run_guardrail(
|
||||
data=data_other, event_type=GuardrailEventHooks.pre_call
|
||||
)
|
||||
is True
|
||||
)
|
||||
assert custom_guardrail.should_run_guardrail(data=data_other, event_type=GuardrailEventHooks.pre_call) is True
|
||||
|
||||
# Test 5: empty opt-out list → still runs
|
||||
data_empty = {
|
||||
|
|
@ -475,12 +422,7 @@ class TestCustomGuardrailShouldRunGuardrail:
|
|||
"messages": [{"role": "user", "content": "test"}],
|
||||
"metadata": {"opted_out_global_guardrails": []},
|
||||
}
|
||||
assert (
|
||||
custom_guardrail.should_run_guardrail(
|
||||
data=data_empty, event_type=GuardrailEventHooks.pre_call
|
||||
)
|
||||
is True
|
||||
)
|
||||
assert custom_guardrail.should_run_guardrail(data=data_empty, event_type=GuardrailEventHooks.pre_call) is True
|
||||
|
||||
# Test 6: malformed value (bool instead of list) → safely ignored, guardrail runs
|
||||
data_malformed = {
|
||||
|
|
@ -489,10 +431,7 @@ class TestCustomGuardrailShouldRunGuardrail:
|
|||
"metadata": {"opted_out_global_guardrails": True},
|
||||
}
|
||||
assert (
|
||||
custom_guardrail.should_run_guardrail(
|
||||
data=data_malformed, event_type=GuardrailEventHooks.pre_call
|
||||
)
|
||||
is True
|
||||
custom_guardrail.should_run_guardrail(data=data_malformed, event_type=GuardrailEventHooks.pre_call) is True
|
||||
)
|
||||
|
||||
def test_should_run_guardrail_opt_out_does_not_affect_non_global(self):
|
||||
|
|
@ -515,12 +454,69 @@ class TestCustomGuardrailShouldRunGuardrail:
|
|||
"guardrails": ["opt_in_guardrail"],
|
||||
},
|
||||
}
|
||||
assert (
|
||||
non_global.should_run_guardrail(
|
||||
data=data, event_type=GuardrailEventHooks.pre_call
|
||||
)
|
||||
is True
|
||||
assert non_global.should_run_guardrail(data=data, event_type=GuardrailEventHooks.pre_call) is True
|
||||
|
||||
def test_should_run_guardrail_suppressed_by_auto_router_compression(self):
|
||||
"""An auto router's own compression policy can suppress an otherwise-eligible
|
||||
guardrail, even one that is default_on and explicitly requested."""
|
||||
from litellm.proxy.guardrails import auto_router_compression
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
|
||||
always_on = CustomGuardrail(
|
||||
guardrail_name="headroom-default",
|
||||
default_on=True,
|
||||
event_hook=GuardrailEventHooks.pre_call,
|
||||
)
|
||||
token = auto_router_compression._suppressed_compression_guardrails.set(frozenset({"headroom-default"}))
|
||||
try:
|
||||
assert (
|
||||
always_on.should_run_guardrail(data={"model": "smart-router"}, event_type=GuardrailEventHooks.pre_call)
|
||||
is False
|
||||
)
|
||||
finally:
|
||||
auto_router_compression._suppressed_compression_guardrails.reset(token)
|
||||
|
||||
def test_should_run_guardrail_suppression_does_not_affect_other_names(self):
|
||||
from litellm.proxy.guardrails import auto_router_compression
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
|
||||
always_on = CustomGuardrail(
|
||||
guardrail_name="headroom-default",
|
||||
default_on=True,
|
||||
event_hook=GuardrailEventHooks.pre_call,
|
||||
)
|
||||
token = auto_router_compression._suppressed_compression_guardrails.set(frozenset({"some-other-guardrail"}))
|
||||
try:
|
||||
assert (
|
||||
always_on.should_run_guardrail(data={"model": "smart-router"}, event_type=GuardrailEventHooks.pre_call)
|
||||
is True
|
||||
)
|
||||
finally:
|
||||
auto_router_compression._suppressed_compression_guardrails.reset(token)
|
||||
|
||||
def test_request_metadata_can_never_suppress_a_guardrail(self):
|
||||
"""Regression (security): suppression state is request-scoped and server-set,
|
||||
never read from metadata. Metadata reaches spend logs the caller can read, so
|
||||
anything honored from there is something a later request could replay to switch
|
||||
off a PII or content-filter guardrail for itself."""
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
|
||||
always_on = CustomGuardrail(
|
||||
guardrail_name="headroom-default",
|
||||
default_on=True,
|
||||
event_hook=GuardrailEventHooks.pre_call,
|
||||
)
|
||||
forged = {
|
||||
"model": "smart-router",
|
||||
"metadata": {
|
||||
"_auto_router_suppressed_compression_guardrails": [
|
||||
"headroom-default",
|
||||
"any-token:headroom-default",
|
||||
],
|
||||
},
|
||||
}
|
||||
|
||||
assert always_on.should_run_guardrail(data=forged, event_type=GuardrailEventHooks.pre_call) is True
|
||||
|
||||
|
||||
class TestApplyGuardrailCheck:
|
||||
|
|
@ -559,35 +555,33 @@ class TestApplyGuardrailCheck:
|
|||
child_with_override = ChildGuardrailWithOverride()
|
||||
|
||||
# Test: CustomGuardrail itself has apply_guardrail in its __dict__
|
||||
assert (
|
||||
"apply_guardrail" in type(CustomGuardrail()).__dict__
|
||||
), "CustomGuardrail should have apply_guardrail in its own __dict__"
|
||||
assert "apply_guardrail" in type(CustomGuardrail()).__dict__, (
|
||||
"CustomGuardrail should have apply_guardrail in its own __dict__"
|
||||
)
|
||||
|
||||
# Test: ParentGuardrail inherits but doesn't override, so it should NOT be in __dict__
|
||||
assert (
|
||||
"apply_guardrail" not in type(parent_instance).__dict__
|
||||
), "ParentGuardrail should NOT have apply_guardrail in its own __dict__ (only inherited)"
|
||||
assert "apply_guardrail" not in type(parent_instance).__dict__, (
|
||||
"ParentGuardrail should NOT have apply_guardrail in its own __dict__ (only inherited)"
|
||||
)
|
||||
|
||||
# Test: ChildGuardrailWithoutOverride only inherits, should NOT be in __dict__
|
||||
assert (
|
||||
"apply_guardrail" not in type(child_without_override).__dict__
|
||||
), "ChildGuardrailWithoutOverride should NOT have apply_guardrail in its own __dict__ (only inherited)"
|
||||
assert "apply_guardrail" not in type(child_without_override).__dict__, (
|
||||
"ChildGuardrailWithoutOverride should NOT have apply_guardrail in its own __dict__ (only inherited)"
|
||||
)
|
||||
|
||||
# Test: ChildGuardrailWithOverride overrides the method, SHOULD be in __dict__
|
||||
assert (
|
||||
"apply_guardrail" in type(child_with_override).__dict__
|
||||
), "ChildGuardrailWithOverride SHOULD have apply_guardrail in its own __dict__ (overridden)"
|
||||
assert "apply_guardrail" in type(child_with_override).__dict__, (
|
||||
"ChildGuardrailWithOverride SHOULD have apply_guardrail in its own __dict__ (overridden)"
|
||||
)
|
||||
|
||||
# Verify that all instances still have the method via inheritance (hasattr)
|
||||
assert hasattr(
|
||||
parent_instance, "apply_guardrail"
|
||||
), "All instances should have apply_guardrail via inheritance"
|
||||
assert hasattr(
|
||||
child_without_override, "apply_guardrail"
|
||||
), "All instances should have apply_guardrail via inheritance"
|
||||
assert hasattr(
|
||||
child_with_override, "apply_guardrail"
|
||||
), "All instances should have apply_guardrail via inheritance"
|
||||
assert hasattr(parent_instance, "apply_guardrail"), "All instances should have apply_guardrail via inheritance"
|
||||
assert hasattr(child_without_override, "apply_guardrail"), (
|
||||
"All instances should have apply_guardrail via inheritance"
|
||||
)
|
||||
assert hasattr(child_with_override, "apply_guardrail"), (
|
||||
"All instances should have apply_guardrail via inheritance"
|
||||
)
|
||||
|
||||
|
||||
class TestGuardrailLoggingAggregation:
|
||||
|
|
@ -614,11 +608,7 @@ class TestGuardrailLoggingAggregation:
|
|||
|
||||
def test_appends_to_existing_metadata_list(self):
|
||||
request_data = {
|
||||
"metadata": {
|
||||
"standard_logging_guardrail_information": [
|
||||
{"guardrail_name": "existing_guardrail"}
|
||||
]
|
||||
}
|
||||
"metadata": {"standard_logging_guardrail_information": [{"guardrail_name": "existing_guardrail"}]}
|
||||
}
|
||||
|
||||
self._invoke_add_log(request_data)
|
||||
|
|
@ -630,11 +620,7 @@ class TestGuardrailLoggingAggregation:
|
|||
assert info[1]["guardrail_name"] == "test_guardrail"
|
||||
|
||||
def test_converts_existing_metadata_dict_to_list(self):
|
||||
request_data = {
|
||||
"metadata": {
|
||||
"standard_logging_guardrail_information": {"guardrail_name": "legacy"}
|
||||
}
|
||||
}
|
||||
request_data = {"metadata": {"standard_logging_guardrail_information": {"guardrail_name": "legacy"}}}
|
||||
|
||||
self._invoke_add_log(request_data)
|
||||
|
||||
|
|
@ -646,18 +632,12 @@ class TestGuardrailLoggingAggregation:
|
|||
|
||||
def test_appends_to_litellm_metadata(self):
|
||||
request_data = {
|
||||
"litellm_metadata": {
|
||||
"standard_logging_guardrail_information": [
|
||||
{"guardrail_name": "litellm_existing"}
|
||||
]
|
||||
}
|
||||
"litellm_metadata": {"standard_logging_guardrail_information": [{"guardrail_name": "litellm_existing"}]}
|
||||
}
|
||||
|
||||
self._invoke_add_log(request_data)
|
||||
|
||||
info = request_data["litellm_metadata"][
|
||||
"standard_logging_guardrail_information"
|
||||
]
|
||||
info = request_data["litellm_metadata"]["standard_logging_guardrail_information"]
|
||||
assert isinstance(info, list)
|
||||
assert len(info) == 2
|
||||
assert info[1]["guardrail_name"] == "test_guardrail"
|
||||
|
|
@ -674,12 +654,10 @@ class TestGuardrailLoggingAggregation:
|
|||
|
||||
self._invoke_add_log(request_data)
|
||||
|
||||
assert (
|
||||
"standard_logging_guardrail_information" not in request_data["metadata"]
|
||||
), "entry landed in the caller's metadata, where the spend log does not read it"
|
||||
info = request_data["litellm_metadata"][
|
||||
"standard_logging_guardrail_information"
|
||||
]
|
||||
assert "standard_logging_guardrail_information" not in request_data["metadata"], (
|
||||
"entry landed in the caller's metadata, where the spend log does not read it"
|
||||
)
|
||||
info = request_data["litellm_metadata"]["standard_logging_guardrail_information"]
|
||||
assert len(info) == 1
|
||||
assert info[0]["guardrail_name"] == "test_guardrail"
|
||||
|
||||
|
|
@ -697,9 +675,7 @@ class TestGuardrailLoggingAggregation:
|
|||
}
|
||||
|
||||
self._invoke_add_log(request_data)
|
||||
add_guardrail_to_applied_guardrails_header(
|
||||
request_data=request_data, guardrail_name="test_guardrail"
|
||||
)
|
||||
add_guardrail_to_applied_guardrails_header(request_data=request_data, guardrail_name="test_guardrail")
|
||||
|
||||
buckets = {
|
||||
key
|
||||
|
|
@ -745,9 +721,7 @@ class TestGuardrailOtelSpanEmission:
|
|||
|
||||
assert len(captured) == 1
|
||||
emitted = captured[0]
|
||||
recorded = request_data["metadata"]["standard_logging_guardrail_information"][
|
||||
-1
|
||||
]
|
||||
recorded = request_data["metadata"]["standard_logging_guardrail_information"][-1]
|
||||
assert emitted is recorded
|
||||
assert emitted["guardrail_name"] == "emit_guard"
|
||||
assert emitted["start_time"] == 1.0
|
||||
|
|
@ -757,9 +731,7 @@ class TestGuardrailOtelSpanEmission:
|
|||
def _boom(_entry):
|
||||
raise RuntimeError("otel exporter down")
|
||||
|
||||
monkeypatch.setattr(
|
||||
"litellm.integrations.otel.logger.emit_guardrail_span", _boom
|
||||
)
|
||||
monkeypatch.setattr("litellm.integrations.otel.logger.emit_guardrail_span", _boom)
|
||||
|
||||
request_data = {"metadata": {}}
|
||||
self._record(self._make_guardrail(), request_data)
|
||||
|
|
@ -856,9 +828,7 @@ class TestGuardrailSensitiveFieldStripping:
|
|||
duration=1.0,
|
||||
)
|
||||
|
||||
logged_response = request_data["metadata"][
|
||||
"standard_logging_guardrail_information"
|
||||
][0]["guardrail_response"]
|
||||
logged_response = request_data["metadata"]["standard_logging_guardrail_information"][0]["guardrail_response"]
|
||||
assert "secret_fields" not in logged_response
|
||||
assert "sk-live-SHOULD-NOT-APPEAR" not in json.dumps(logged_response)
|
||||
|
||||
|
|
@ -871,9 +841,7 @@ class TestGuardrailSensitiveFieldStripping:
|
|||
guardrail_json_response=[
|
||||
{
|
||||
"result": "ok",
|
||||
"secret_fields": {
|
||||
"raw_headers": {"authorization": "Bearer sk-secret"}
|
||||
},
|
||||
"secret_fields": {"raw_headers": {"authorization": "Bearer sk-secret"}},
|
||||
},
|
||||
{"result": "also_ok"},
|
||||
],
|
||||
|
|
@ -927,9 +895,7 @@ class TestGuardrailResponseCredentialMasking:
|
|||
duration=1.0,
|
||||
)
|
||||
|
||||
logged = request_data["metadata"]["standard_logging_guardrail_information"][0][
|
||||
"guardrail_response"
|
||||
]
|
||||
logged = request_data["metadata"]["standard_logging_guardrail_information"][0]["guardrail_response"]
|
||||
|
||||
masked_key = logged["metadata_snapshot"]["callback_vars"]["langsmith_api_key"]
|
||||
assert masked_key != plaintext_key
|
||||
|
|
@ -938,10 +904,7 @@ class TestGuardrailResponseCredentialMasking:
|
|||
|
||||
assert logged["model"] == "gpt-4o-mini"
|
||||
assert logged["messages"] == [{"role": "user", "content": "hi"}]
|
||||
assert (
|
||||
logged["metadata_snapshot"]["callback_vars"]["langsmith_project"]
|
||||
== "proj-name"
|
||||
)
|
||||
assert logged["metadata_snapshot"]["callback_vars"]["langsmith_project"] == "proj-name"
|
||||
|
||||
def test_nested_user_api_key_auth_metadata_is_masked(self):
|
||||
import json
|
||||
|
|
@ -1000,9 +963,7 @@ class TestGuardrailResponseCredentialMasking:
|
|||
request_data: dict = {"metadata": {}}
|
||||
|
||||
guardrail.add_standard_logging_guardrail_information_to_request_data(
|
||||
guardrail_json_response={
|
||||
"filters": [{"regex": r"\d{3}-\d{2}-\d{4}", "action": "BLOCKED"}]
|
||||
},
|
||||
guardrail_json_response={"filters": [{"regex": r"\d{3}-\d{2}-\d{4}", "action": "BLOCKED"}]},
|
||||
request_data=request_data,
|
||||
guardrail_status="success",
|
||||
)
|
||||
|
|
@ -1025,9 +986,7 @@ class TestGuardrailResponseCredentialMasking:
|
|||
guardrail_status="success",
|
||||
)
|
||||
|
||||
logged = request_data["metadata"]["standard_logging_guardrail_information"][0][
|
||||
"guardrail_response"
|
||||
]
|
||||
logged = request_data["metadata"]["standard_logging_guardrail_information"][0]["guardrail_response"]
|
||||
assert logged["flagged"] is True
|
||||
assert logged["score"] == 0.94
|
||||
assert logged["tokens_used"] == 42
|
||||
|
|
@ -1039,18 +998,14 @@ class TestGuardrailResponseCredentialMasking:
|
|||
plaintext = "lsv2_pt_abcdef1234567890"
|
||||
|
||||
guardrail.add_standard_logging_guardrail_information_to_request_data(
|
||||
guardrail_json_response={
|
||||
"metadata_snapshot": {
|
||||
"callback_vars": {"langsmith_api_key": plaintext}
|
||||
}
|
||||
},
|
||||
guardrail_json_response={"metadata_snapshot": {"callback_vars": {"langsmith_api_key": plaintext}}},
|
||||
request_data=request_data,
|
||||
guardrail_status="success",
|
||||
)
|
||||
|
||||
masked = request_data["metadata"]["standard_logging_guardrail_information"][0][
|
||||
"guardrail_response"
|
||||
]["metadata_snapshot"]["callback_vars"]["langsmith_api_key"]
|
||||
masked = request_data["metadata"]["standard_logging_guardrail_information"][0]["guardrail_response"][
|
||||
"metadata_snapshot"
|
||||
]["callback_vars"]["langsmith_api_key"]
|
||||
assert masked != plaintext
|
||||
assert masked.startswith(plaintext[:4])
|
||||
assert masked.endswith(plaintext[-4:])
|
||||
|
|
@ -1544,9 +1499,7 @@ class TestEventTypeLogging:
|
|||
guardrail = TestGuardrail()
|
||||
request_data = {"metadata": {}}
|
||||
|
||||
await guardrail.apply_guardrail(
|
||||
inputs={"texts": ["x"]}, request_data=request_data
|
||||
)
|
||||
await guardrail.apply_guardrail(inputs={"texts": ["x"]}, request_data=request_data)
|
||||
|
||||
logged_info = request_data["metadata"]["standard_logging_guardrail_information"]
|
||||
assert len(logged_info) == 1, (
|
||||
|
|
@ -1588,9 +1541,7 @@ class TestEventTypeLogging:
|
|||
request_data = {"metadata": {}}
|
||||
|
||||
with pytest.raises(ValueError, match="blocked"):
|
||||
await guardrail.apply_guardrail(
|
||||
inputs={"texts": ["x"]}, request_data=request_data
|
||||
)
|
||||
await guardrail.apply_guardrail(inputs={"texts": ["x"]}, request_data=request_data)
|
||||
|
||||
logged_info = request_data["metadata"]["standard_logging_guardrail_information"]
|
||||
assert len(logged_info) == 1
|
||||
|
|
@ -1719,9 +1670,7 @@ class TestTracingFieldsPopulation:
|
|||
guardrail_json_response="blocked",
|
||||
request_data=request_data,
|
||||
guardrail_status="guardrail_intervened",
|
||||
tracing_detail=GuardrailTracingDetail(
|
||||
policy_template="EU AI Act Article 5"
|
||||
),
|
||||
tracing_detail=GuardrailTracingDetail(policy_template="EU AI Act Article 5"),
|
||||
)
|
||||
|
||||
slg_list = request_data["metadata"]["standard_logging_guardrail_information"]
|
||||
|
|
@ -1763,13 +1712,7 @@ class TestCustomGuardrailSpendLogMatchRedaction:
|
|||
cg = CustomGuardrail(guardrail_name="test-rail")
|
||||
raw = {
|
||||
"assessments": [
|
||||
{
|
||||
"sensitiveInformationPolicy": {
|
||||
"piiEntities": [
|
||||
{"type": "NAME", "match": "GG", "action": "BLOCKED"}
|
||||
]
|
||||
}
|
||||
}
|
||||
{"sensitiveInformationPolicy": {"piiEntities": [{"type": "NAME", "match": "GG", "action": "BLOCKED"}]}}
|
||||
]
|
||||
}
|
||||
request_data: dict = {"metadata": {}}
|
||||
|
|
@ -1780,17 +1723,10 @@ class TestCustomGuardrailSpendLogMatchRedaction:
|
|||
)
|
||||
slg = request_data["metadata"]["standard_logging_guardrail_information"][0]
|
||||
assert (
|
||||
slg["guardrail_response"]["assessments"][0]["sensitiveInformationPolicy"][
|
||||
"piiEntities"
|
||||
][0]["match"]
|
||||
slg["guardrail_response"]["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"][0]["match"]
|
||||
== "[REDACTED]"
|
||||
)
|
||||
assert (
|
||||
raw["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"][0][
|
||||
"match"
|
||||
]
|
||||
== "GG"
|
||||
)
|
||||
assert raw["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"][0]["match"] == "GG"
|
||||
|
||||
def test_add_standard_logging_redacts_regex_field(self):
|
||||
cg = CustomGuardrail(guardrail_name="test-rail")
|
||||
|
|
|
|||
|
|
@ -0,0 +1,381 @@
|
|||
"""Unit tests for litellm.proxy.guardrails.auto_router_compression."""
|
||||
|
||||
import json
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.integrations.custom_guardrail import CustomGuardrail
|
||||
from litellm.proxy.guardrails import auto_router_compression
|
||||
from litellm.proxy.guardrails.auto_router_compression import (
|
||||
AutoRouterCompressionPolicy,
|
||||
arm_pre_call,
|
||||
messages_for_routing,
|
||||
policy_for_model,
|
||||
policy_from_litellm_params,
|
||||
)
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
from litellm.types.utils import GenericGuardrailAPIInputs
|
||||
|
||||
|
||||
class TestPolicyFromLitellmParams:
|
||||
def test_neither_key_set_is_no_policy(self):
|
||||
assert policy_from_litellm_params({}) is None
|
||||
|
||||
def test_routing_only(self):
|
||||
policy = policy_from_litellm_params({"auto_router_routing_compression": "headroom-a"})
|
||||
assert policy == AutoRouterCompressionPolicy(routing="headroom-a", model=None)
|
||||
|
||||
def test_none_sentinel_normalizes_to_no_compression(self):
|
||||
policy = policy_from_litellm_params(
|
||||
{"auto_router_routing_compression": "headroom-a", "auto_router_model_compression": "none"}
|
||||
)
|
||||
assert policy == AutoRouterCompressionPolicy(routing="headroom-a", model=None)
|
||||
|
||||
def test_none_sentinel_is_case_insensitive(self):
|
||||
policy = policy_from_litellm_params({"auto_router_routing_compression": "NONE"})
|
||||
assert policy == AutoRouterCompressionPolicy(routing=None, model=None)
|
||||
|
||||
def test_is_same_true_for_matching_names(self):
|
||||
policy = policy_from_litellm_params(
|
||||
{"auto_router_routing_compression": "x", "auto_router_model_compression": "x"}
|
||||
)
|
||||
assert policy.is_same is True
|
||||
|
||||
def test_is_same_false_for_different_names(self):
|
||||
policy = policy_from_litellm_params(
|
||||
{"auto_router_routing_compression": "x", "auto_router_model_compression": "y"}
|
||||
)
|
||||
assert policy.is_same is False
|
||||
|
||||
def test_is_same_true_when_both_no_compression(self):
|
||||
policy = policy_from_litellm_params(
|
||||
{"auto_router_routing_compression": "none", "auto_router_model_compression": "none"}
|
||||
)
|
||||
assert policy.is_same is True
|
||||
|
||||
|
||||
class _FakeRouter:
|
||||
"""Minimal stand-in for litellm.Router.get_model_list, for policy_for_model."""
|
||||
|
||||
def __init__(self, deployments: list[dict[str, Any]]):
|
||||
self._deployments = deployments
|
||||
|
||||
def get_model_list(self, model_name, team_id=None):
|
||||
return [d for d in self._deployments if d.get("model_name") == model_name]
|
||||
|
||||
|
||||
def _marker(compression: dict[str, str], tags: list[str] | None = None) -> dict[str, Any]:
|
||||
return {
|
||||
"model_name": "smart-router",
|
||||
"litellm_params": {
|
||||
"model": "auto_router/complexity_router",
|
||||
**compression,
|
||||
**({"tags": tags} if tags is not None else {}),
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
class TestPolicyForModel:
|
||||
def test_no_router_returns_none(self):
|
||||
assert policy_for_model(llm_router=None, model_alias="smart-router", team_id=None, request_tags=()) is None
|
||||
|
||||
def test_no_marker_deployment_returns_none(self):
|
||||
router = _FakeRouter([{"model_name": "smart-router", "litellm_params": {"model": "openai/gpt-4o-mini"}}])
|
||||
assert policy_for_model(llm_router=router, model_alias="smart-router", team_id=None, request_tags=()) is None
|
||||
|
||||
def test_marker_deployment_without_policy_returns_none(self):
|
||||
router = _FakeRouter(
|
||||
[{"model_name": "smart-router", "litellm_params": {"model": "auto_router/complexity_router"}}]
|
||||
)
|
||||
assert policy_for_model(llm_router=router, model_alias="smart-router", team_id=None, request_tags=()) is None
|
||||
|
||||
def test_marker_deployment_with_policy_is_found(self):
|
||||
router = _FakeRouter(
|
||||
[_marker({"auto_router_routing_compression": "headroom-a", "auto_router_model_compression": "none"})]
|
||||
)
|
||||
policy = policy_for_model(llm_router=router, model_alias="smart-router", team_id=None, request_tags=())
|
||||
assert policy == AutoRouterCompressionPolicy(routing="headroom-a", model=None)
|
||||
|
||||
def test_picks_the_marker_whose_tags_the_request_carries(self):
|
||||
"""Regression: an alias with several tag-scoped markers must not suppress one
|
||||
marker's guardrail and then route under a different marker's policy."""
|
||||
router = _FakeRouter(
|
||||
[
|
||||
_marker({"auto_router_routing_compression": "headroom-eu"}, tags=["eu"]),
|
||||
_marker({"auto_router_routing_compression": "headroom-us"}, tags=["us"]),
|
||||
]
|
||||
)
|
||||
|
||||
eu = policy_for_model(llm_router=router, model_alias="smart-router", team_id=None, request_tags=("eu",))
|
||||
us = policy_for_model(llm_router=router, model_alias="smart-router", team_id=None, request_tags=("us",))
|
||||
|
||||
assert eu == AutoRouterCompressionPolicy(routing="headroom-eu", model=None)
|
||||
assert us == AutoRouterCompressionPolicy(routing="headroom-us", model=None)
|
||||
|
||||
def test_untagged_marker_matches_any_request(self):
|
||||
router = _FakeRouter([_marker({"auto_router_routing_compression": "headroom-a"})])
|
||||
policy = policy_for_model(
|
||||
llm_router=router, model_alias="smart-router", team_id=None, request_tags=("anything",)
|
||||
)
|
||||
assert policy == AutoRouterCompressionPolicy(routing="headroom-a", model=None)
|
||||
|
||||
def test_a_marker_scoped_to_other_tags_is_never_the_fallback(self):
|
||||
"""Regression: a "us" request must not fall back to an "eu" marker's policy."""
|
||||
router = _FakeRouter(
|
||||
[
|
||||
_marker({"auto_router_routing_compression": "headroom-eu"}, tags=["eu"]),
|
||||
_marker({"auto_router_routing_compression": "headroom-default"}),
|
||||
]
|
||||
)
|
||||
policy = policy_for_model(llm_router=router, model_alias="smart-router", team_id=None, request_tags=("us",))
|
||||
assert policy == AutoRouterCompressionPolicy(routing="headroom-default", model=None)
|
||||
|
||||
def test_no_untagged_fallback_means_no_policy(self):
|
||||
"""No matching marker means no policy, not an unrelated slice's compression."""
|
||||
router = _FakeRouter([_marker({"auto_router_routing_compression": "headroom-eu"}, tags=["eu"])])
|
||||
assert (
|
||||
policy_for_model(llm_router=router, model_alias="smart-router", team_id=None, request_tags=("us",)) is None
|
||||
)
|
||||
|
||||
def test_tag_scoped_marker_takes_precedence_over_untagged(self):
|
||||
"""Regression: when multiple markers exist, the tag-scoped one the request
|
||||
actually matches should be used, not the first untagged one."""
|
||||
router = _FakeRouter(
|
||||
[
|
||||
_marker({"auto_router_routing_compression": "headroom-untagged"}),
|
||||
_marker({"auto_router_routing_compression": "headroom-eu"}, tags=["eu"]),
|
||||
]
|
||||
)
|
||||
policy = policy_for_model(llm_router=router, model_alias="smart-router", team_id=None, request_tags=("eu",))
|
||||
assert policy == AutoRouterCompressionPolicy(routing="headroom-eu", model=None)
|
||||
|
||||
|
||||
class _RecordingCompressionGuardrail(CustomGuardrail):
|
||||
"""A guardrail whose apply_guardrail marks every text message as compressed."""
|
||||
|
||||
def __init__(self, guardrail_name: str):
|
||||
super().__init__(guardrail_name=guardrail_name)
|
||||
self.request_data_seen: list[dict] = []
|
||||
|
||||
async def apply_guardrail(
|
||||
self, inputs: GenericGuardrailAPIInputs, request_data: dict, input_type: str, logging_obj=None
|
||||
) -> GenericGuardrailAPIInputs:
|
||||
self.request_data_seen.append(request_data)
|
||||
structured_messages = inputs.get("structured_messages") or []
|
||||
compressed = [{**m, "content": f"[COMPRESSED] {m.get('content')}"} for m in structured_messages]
|
||||
return {**inputs, "structured_messages": compressed}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def registered_guardrail(monkeypatch):
|
||||
import litellm
|
||||
from litellm.proxy.guardrails import guardrail_registry
|
||||
|
||||
# Registered under a compression provider name: both hops refuse a name that does
|
||||
# not resolve to one, so a bare callback would (correctly) never be used.
|
||||
monkeypatch.setitem(guardrail_registry.guardrail_class_registry, "headroom", _RecordingCompressionGuardrail)
|
||||
guardrail = _RecordingCompressionGuardrail(guardrail_name="fake-compress")
|
||||
litellm.logging_callback_manager.add_litellm_callback(guardrail)
|
||||
yield guardrail
|
||||
litellm.logging_callback_manager.remove_callback_from_all_lists(guardrail)
|
||||
|
||||
|
||||
class _NonCompressionGuardrail(CustomGuardrail):
|
||||
"""A guardrail that is not a compression provider, e.g. a PII or content filter."""
|
||||
|
||||
def __init__(self, guardrail_name: str):
|
||||
super().__init__(guardrail_name=guardrail_name)
|
||||
self.called = False
|
||||
|
||||
async def apply_guardrail(
|
||||
self, inputs: GenericGuardrailAPIInputs, request_data: dict, input_type: str, logging_obj=None
|
||||
) -> GenericGuardrailAPIInputs:
|
||||
self.called = True
|
||||
return inputs
|
||||
|
||||
|
||||
class TestArmPreCall:
|
||||
@pytest.mark.asyncio
|
||||
async def test_no_router_is_noop(self):
|
||||
data = {"model": "smart-router", "messages": [{"role": "user", "content": "hi"}]}
|
||||
await arm_pre_call(data=data, llm_router=None)
|
||||
assert "metadata" not in data
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_no_policy_does_not_create_metadata_bucket(self):
|
||||
router = _FakeRouter([{"model_name": "smart-router", "litellm_params": {"model": "openai/gpt-4o-mini"}}])
|
||||
data = {"model": "smart-router", "messages": [{"role": "user", "content": "hi"}]}
|
||||
await arm_pre_call(data=data, llm_router=router)
|
||||
assert "metadata" not in data
|
||||
assert "litellm_metadata" not in data
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_policy_suppresses_active_compression_guardrails(self, monkeypatch):
|
||||
from litellm.proxy.guardrails import guardrail_registry
|
||||
|
||||
monkeypatch.setitem(
|
||||
guardrail_registry.guardrail_class_registry, "fake-provider", _RecordingCompressionGuardrail
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.guardrails.auto_router_compression.COMPRESSION_GUARDRAIL_PROVIDERS",
|
||||
frozenset({"fake-provider"}),
|
||||
)
|
||||
import litellm
|
||||
|
||||
always_on = _RecordingCompressionGuardrail(guardrail_name="always-on-compression")
|
||||
litellm.logging_callback_manager.add_litellm_callback(always_on)
|
||||
try:
|
||||
router = _FakeRouter(
|
||||
[
|
||||
{
|
||||
"model_name": "smart-router",
|
||||
"litellm_params": {
|
||||
"model": "auto_router/complexity_router",
|
||||
"auto_router_routing_compression": "headroom-a",
|
||||
"auto_router_model_compression": "none",
|
||||
},
|
||||
}
|
||||
]
|
||||
)
|
||||
data = {"model": "smart-router", "messages": [{"role": "user", "content": "hi"}]}
|
||||
await arm_pre_call(data=data, llm_router=router)
|
||||
assert auto_router_compression.suppressed_compression_guardrails() == frozenset({"always-on-compression"})
|
||||
# Suppression state must never ride along in metadata: that reaches spend
|
||||
# logs the caller can read, and anything there is replayable.
|
||||
assert "always-on-compression" not in json.dumps(data.get("metadata", {}))
|
||||
assert always_on.should_run_guardrail(data=data, event_type=GuardrailEventHooks.pre_call) is False
|
||||
finally:
|
||||
litellm.logging_callback_manager.remove_callback_from_all_lists(always_on)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_suppression_state_never_enters_request_metadata(self):
|
||||
"""Regression (security): metadata reaches spend logs, so a suppression list
|
||||
there is one a caller could read back and replay to disable a guardrail."""
|
||||
guardrail = _RecordingCompressionGuardrail(guardrail_name="always-on-compression")
|
||||
import litellm
|
||||
|
||||
litellm.logging_callback_manager.add_litellm_callback(guardrail)
|
||||
try:
|
||||
router = _FakeRouter([_marker({"auto_router_routing_compression": "headroom-a"})])
|
||||
data = {"model": "smart-router", "messages": [{"role": "user", "content": "hi"}]}
|
||||
await arm_pre_call(data=data, llm_router=router)
|
||||
assert "suppress" not in json.dumps(data).lower()
|
||||
finally:
|
||||
litellm.logging_callback_manager.remove_callback_from_all_lists(guardrail)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_model_side_guardrail_is_requested_even_when_not_default_on(self, monkeypatch):
|
||||
import litellm
|
||||
from litellm.proxy.guardrails import guardrail_registry
|
||||
|
||||
monkeypatch.setitem(guardrail_registry.guardrail_class_registry, "headroom", _RecordingCompressionGuardrail)
|
||||
active = _RecordingCompressionGuardrail(guardrail_name="headroom-b")
|
||||
litellm.logging_callback_manager.add_litellm_callback(active)
|
||||
router = _FakeRouter(
|
||||
[
|
||||
{
|
||||
"model_name": "smart-router",
|
||||
"litellm_params": {
|
||||
"model": "auto_router/complexity_router",
|
||||
"auto_router_routing_compression": "none",
|
||||
"auto_router_model_compression": "headroom-b",
|
||||
},
|
||||
}
|
||||
]
|
||||
)
|
||||
data = {"model": "smart-router", "messages": [{"role": "user", "content": "hi"}]}
|
||||
try:
|
||||
await arm_pre_call(data=data, llm_router=router)
|
||||
assert data["metadata"]["guardrails"] == ["headroom-b"]
|
||||
finally:
|
||||
litellm.logging_callback_manager.remove_callback_from_all_lists(active)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_arm_pre_call_keeps_no_copy_of_the_prompt(self):
|
||||
"""Regression (security): arm_pre_call runs before the guardrails, so any copy it
|
||||
kept would be pre-masking text that routing then POSTs to an external service."""
|
||||
router = _FakeRouter([_marker({"auto_router_routing_compression": "headroom-a"})])
|
||||
data = {"model": "smart-router", "messages": [{"role": "user", "content": "my ssn is 123-45-6789"}]}
|
||||
|
||||
await arm_pre_call(data=data, llm_router=router)
|
||||
|
||||
assert "123-45-6789" not in json.dumps(data.get("metadata", {}))
|
||||
assert not hasattr(auto_router_compression, "_routing_messages_snapshot")
|
||||
|
||||
|
||||
class TestMessagesForRouting:
|
||||
@pytest.mark.asyncio
|
||||
async def test_no_policy_returns_none(self):
|
||||
assert await messages_for_routing(policy=None, messages=[], request_kwargs={}) is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_routing_none_with_no_model_compression_returns_none(self):
|
||||
"""Nothing compressed either hop, so the caller's own messages are already right."""
|
||||
policy = AutoRouterCompressionPolicy(routing=None, model=None)
|
||||
assert await messages_for_routing(policy=policy, messages=[], request_kwargs={}) is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_routing_none_never_reaches_for_a_pre_guardrail_copy(self):
|
||||
"""No uncompressed copy survives the model hop, and keeping one would mean
|
||||
retaining the pre-masking text. Routing reads what it has."""
|
||||
policy = AutoRouterCompressionPolicy(routing=None, model="headroom-a")
|
||||
model_compressed = [{"role": "user", "content": "[COMPRESSED] the full original conversation"}]
|
||||
|
||||
assert await messages_for_routing(policy=policy, messages=model_compressed, request_kwargs={}) is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_unknown_guardrail_name_routes_on_the_uncompressed_messages(self):
|
||||
policy = AutoRouterCompressionPolicy(routing="does-not-exist", model=None)
|
||||
messages = [{"role": "user", "content": "hi"}]
|
||||
result = await messages_for_routing(policy=policy, messages=messages, request_kwargs={})
|
||||
assert result == messages
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_compresses_via_the_named_guardrail(self, registered_guardrail):
|
||||
policy = AutoRouterCompressionPolicy(routing="fake-compress", model=None)
|
||||
messages = [{"role": "user", "content": "hello world"}]
|
||||
result = await messages_for_routing(policy=policy, messages=messages, request_kwargs={})
|
||||
assert result == [{"role": "user", "content": "[COMPRESSED] hello world"}]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_routing_compresses_what_the_other_guardrails_left_behind(self, registered_guardrail):
|
||||
"""Regression (security): routing POSTs its input out, so it must read what the
|
||||
earlier guardrails left behind, not a pre-masking copy."""
|
||||
policy = AutoRouterCompressionPolicy(routing="fake-compress", model="headroom-b")
|
||||
masked = [{"role": "user", "content": "my ssn is [REDACTED]"}]
|
||||
|
||||
result = await messages_for_routing(policy=policy, messages=masked, request_kwargs={})
|
||||
|
||||
assert result == [{"role": "user", "content": "[COMPRESSED] my ssn is [REDACTED]"}]
|
||||
assert registered_guardrail.request_data_seen[0]["messages"] == masked
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_non_compression_guardrail_is_never_invoked_for_routing(self, monkeypatch):
|
||||
"""Regression (security): naming an ordinary guardrail must not turn the routing
|
||||
hop into a way to ship prompts to whatever service backs it."""
|
||||
import litellm
|
||||
|
||||
other = _NonCompressionGuardrail(guardrail_name="pii-filter")
|
||||
litellm.logging_callback_manager.add_litellm_callback(other)
|
||||
try:
|
||||
policy = AutoRouterCompressionPolicy(routing="pii-filter", model=None)
|
||||
messages = [{"role": "user", "content": "my ssn is 123-45-6789"}]
|
||||
|
||||
result = await messages_for_routing(policy=policy, messages=messages, request_kwargs={})
|
||||
|
||||
assert other.called is False
|
||||
assert result == messages
|
||||
finally:
|
||||
litellm.logging_callback_manager.remove_callback_from_all_lists(other)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_guardrail_receives_a_throwaway_request_data_not_the_real_request_kwargs(self, registered_guardrail):
|
||||
"""Regression: a guardrail writes stats onto the request_data it is given, so
|
||||
passing the caller's own would double-count into extract_compression_saved_tokens."""
|
||||
policy = AutoRouterCompressionPolicy(routing="fake-compress", model=None)
|
||||
messages = [{"role": "user", "content": "hi"}]
|
||||
request_kwargs = {"metadata": {}}
|
||||
await messages_for_routing(policy=policy, messages=messages, request_kwargs=request_kwargs)
|
||||
assert registered_guardrail.request_data_seen[0] is not request_kwargs
|
||||
assert request_kwargs == {"metadata": {}}
|
||||
|
|
@ -376,6 +376,75 @@ class TestProxyBaseLLMRequestProcessing:
|
|||
assert "litellm_logging_obj" not in persisted_body
|
||||
json.dumps(persisted_body)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_common_processing_pre_call_logic_arms_auto_router_compression_before_guardrails(
|
||||
self, monkeypatch
|
||||
):
|
||||
"""arm_pre_call must run before pre_call_hook: an auto router's own compression
|
||||
policy has to be in `data["metadata"]` (naming the model-side guardrail so it
|
||||
runs even if it isn't default_on) by the time guardrails see the request."""
|
||||
from litellm.integrations.custom_guardrail import CustomGuardrail
|
||||
from litellm.proxy.guardrails import guardrail_registry
|
||||
|
||||
# The model hop is only armed for a name that resolves to an active compression
|
||||
# guardrail, so arming it has to have a real one to resolve to.
|
||||
class _FakeCompressionGuardrail(CustomGuardrail):
|
||||
pass
|
||||
|
||||
monkeypatch.setitem(guardrail_registry.guardrail_class_registry, "headroom", _FakeCompressionGuardrail)
|
||||
active_guardrail = _FakeCompressionGuardrail(guardrail_name="headroom-model")
|
||||
litellm.logging_callback_manager.add_litellm_callback(active_guardrail)
|
||||
|
||||
processing_obj = ProxyBaseLLMRequestProcessing(data={})
|
||||
mock_request = MagicMock(spec=Request)
|
||||
mock_request.headers = {}
|
||||
|
||||
async def mock_add_litellm_data_to_request(*args, **kwargs):
|
||||
return {"model": "smart-router", "messages": [{"role": "user", "content": "hi"}]}
|
||||
|
||||
seen_metadata: dict = {}
|
||||
|
||||
async def mock_pre_call_hook(user_api_key_dict, data, call_type):
|
||||
seen_metadata.update(data.get("metadata") or {})
|
||||
return data
|
||||
|
||||
mock_proxy_logging_obj = MagicMock(spec=ProxyLogging)
|
||||
mock_proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=mock_pre_call_hook)
|
||||
monkeypatch.setattr(
|
||||
litellm.proxy.common_request_processing,
|
||||
"add_litellm_data_to_request",
|
||||
mock_add_litellm_data_to_request,
|
||||
)
|
||||
|
||||
fake_llm_router = MagicMock()
|
||||
fake_llm_router.get_model_list.return_value = [
|
||||
{
|
||||
"model_name": "smart-router",
|
||||
"litellm_params": {
|
||||
"model": "auto_router/complexity_router",
|
||||
"auto_router_routing_compression": "none",
|
||||
"auto_router_model_compression": "headroom-model",
|
||||
},
|
||||
}
|
||||
]
|
||||
mock_proxy_config = MagicMock(spec=ProxyConfig)
|
||||
mock_proxy_config._get_hierarchical_router_settings = AsyncMock(return_value=None)
|
||||
|
||||
try:
|
||||
await processing_obj.common_processing_pre_call_logic(
|
||||
request=mock_request,
|
||||
general_settings={},
|
||||
user_api_key_dict=ProxyUserAPIKeyAuth(api_key="sk-test"),
|
||||
proxy_logging_obj=mock_proxy_logging_obj,
|
||||
proxy_config=mock_proxy_config,
|
||||
route_type="acompletion",
|
||||
llm_router=fake_llm_router,
|
||||
)
|
||||
finally:
|
||||
litellm.logging_callback_manager.remove_callback_from_all_lists(active_guardrail)
|
||||
|
||||
assert seen_metadata.get("guardrails") == ["headroom-model"]
|
||||
|
||||
def test_add_dd_apm_tags_for_litellm_call_id_uses_dd_tracing_helper(self, monkeypatch):
|
||||
mock_set_active_span_tag = MagicMock(return_value=True)
|
||||
import litellm.proxy.dd_span_tagger
|
||||
|
|
|
|||
|
|
@ -20,6 +20,7 @@ import respx
|
|||
import litellm
|
||||
from litellm import Router
|
||||
from litellm.exceptions import MidStreamFallbackError
|
||||
from litellm.integrations.custom_guardrail import CustomGuardrail
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
|
||||
from litellm.llms.bedrock.common_utils import BedrockError
|
||||
|
|
@ -9996,6 +9997,215 @@ class TestModelGroupAliasReachesPreRoutingStrategies:
|
|||
)
|
||||
|
||||
|
||||
class TestAutoRouterCompressionDecoupling:
|
||||
"""An auto router's `auto_router_routing_compression` / `auto_router_model_compression`
|
||||
decouple what the routing decision sees from what the model call sees. The one
|
||||
assertion that must hold under any mutation: the strategy can be routed on
|
||||
compressed text while the caller's own `messages` list - the one that would reach
|
||||
the model - is never touched."""
|
||||
|
||||
class _RecordingStrategy:
|
||||
"""Echoes back whatever `messages` it was handed, like every real strategy does."""
|
||||
|
||||
def __init__(self):
|
||||
self.received_messages: list[dict] | None = None
|
||||
|
||||
async def async_pre_routing_hook(
|
||||
self, model, request_kwargs, messages=None, input=None, specific_deployment=False
|
||||
):
|
||||
from litellm.types.router import PreRoutingHookResponse
|
||||
|
||||
self.received_messages = messages
|
||||
return PreRoutingHookResponse(model="gemini-flash", messages=messages)
|
||||
|
||||
class _CompressingGuardrail(CustomGuardrail):
|
||||
def __init__(self, guardrail_name: str):
|
||||
super().__init__(guardrail_name=guardrail_name)
|
||||
self.call_count = 0
|
||||
|
||||
async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None):
|
||||
self.call_count += 1
|
||||
structured_messages = inputs.get("structured_messages") or []
|
||||
compressed = [{**m, "content": f"[COMPRESSED] {m.get('content')}"} for m in structured_messages]
|
||||
return {**inputs, "structured_messages": compressed}
|
||||
|
||||
@staticmethod
|
||||
def _messages() -> list[dict[str, str]]:
|
||||
return [{"role": "user", "content": "What is the capital of France?"}]
|
||||
|
||||
def _router(self, marker_litellm_params: dict) -> tuple[litellm.Router, "_RecordingStrategy"]:
|
||||
from litellm.types.router import TaggedPreRoutingStrategy
|
||||
|
||||
tiers = dict.fromkeys(("SIMPLE", "MEDIUM", "COMPLEX", "REASONING"), "gemini-flash")
|
||||
router = litellm.Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "smart-router",
|
||||
"litellm_params": {
|
||||
"model": "auto_router/complexity_router",
|
||||
"complexity_router_config": {"tiers": tiers},
|
||||
"complexity_router_default_model": "gemini-flash",
|
||||
**marker_litellm_params,
|
||||
},
|
||||
},
|
||||
{
|
||||
"model_name": "gemini-flash",
|
||||
"litellm_params": {"model": "gemini/gemini-3.6-flash", "mock_response": "routed by the tier"},
|
||||
},
|
||||
],
|
||||
)
|
||||
for name in ("auto_routers", "complexity_routers", "adaptive_routers", "quality_routers"):
|
||||
setattr(router, name, {})
|
||||
strategy = self._RecordingStrategy()
|
||||
router.complexity_routers = {"smart-router": [TaggedPreRoutingStrategy(tags=(), strategy=strategy)]}
|
||||
return router, strategy
|
||||
|
||||
@pytest.fixture
|
||||
def registered_guardrail(self, monkeypatch):
|
||||
from litellm.proxy.guardrails import guardrail_registry
|
||||
|
||||
# Registered under a compression provider name: both hops refuse a name that
|
||||
# does not resolve to one, so a bare callback would never be used.
|
||||
monkeypatch.setitem(guardrail_registry.guardrail_class_registry, "headroom", self._CompressingGuardrail)
|
||||
guardrail = self._CompressingGuardrail(guardrail_name="fake-compress")
|
||||
litellm.logging_callback_manager.add_litellm_callback(guardrail)
|
||||
yield guardrail
|
||||
litellm.logging_callback_manager.remove_callback_from_all_lists(guardrail)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_routing_side_compression_never_reaches_the_caller_messages(self, registered_guardrail):
|
||||
router, strategy = self._router(
|
||||
{
|
||||
"auto_router_routing_compression": "fake-compress",
|
||||
"auto_router_model_compression": "none",
|
||||
}
|
||||
)
|
||||
original_messages = self._messages()
|
||||
|
||||
response = await router.async_pre_routing_hook(
|
||||
model="smart-router", request_kwargs={"metadata": {}}, messages=original_messages
|
||||
)
|
||||
|
||||
assert strategy.received_messages == [
|
||||
{"role": "user", "content": "[COMPRESSED] What is the capital of France?"}
|
||||
]
|
||||
assert response.messages == original_messages
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_model_side_compression_alone_leaves_routing_uncompressed(self, registered_guardrail):
|
||||
router, strategy = self._router(
|
||||
{
|
||||
"auto_router_routing_compression": "none",
|
||||
"auto_router_model_compression": "fake-compress",
|
||||
}
|
||||
)
|
||||
original_messages = self._messages()
|
||||
|
||||
response = await router.async_pre_routing_hook(
|
||||
model="smart-router", request_kwargs={"metadata": {}}, messages=original_messages
|
||||
)
|
||||
|
||||
assert strategy.received_messages == original_messages
|
||||
assert response.messages == original_messages
|
||||
assert registered_guardrail.call_count == 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_routing_none_classifies_on_the_live_messages_not_a_pre_guardrail_copy(self, registered_guardrail):
|
||||
"""Routing asked for no compression while the model hop compressed, so the only
|
||||
messages left are that guardrail's output and the strategy classifies on them.
|
||||
|
||||
Keeping a pre-compression copy to classify on instead is what this deliberately
|
||||
gives up: that copy is taken before the pre-call guardrails run, so it still
|
||||
holds whatever a masking guardrail exists to strip, and routing-side compression
|
||||
POSTs its input to an external service."""
|
||||
router, strategy = self._router(
|
||||
{
|
||||
"auto_router_routing_compression": "none",
|
||||
"auto_router_model_compression": "fake-compress",
|
||||
}
|
||||
)
|
||||
model_compressed = [{"role": "user", "content": "[COMPRESSED] What is the capital of France?"}]
|
||||
|
||||
await router.async_pre_routing_hook(
|
||||
model="smart-router", request_kwargs={"metadata": {}}, messages=model_compressed
|
||||
)
|
||||
|
||||
assert strategy.received_messages == model_compressed
|
||||
assert registered_guardrail.call_count == 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.asyncio
|
||||
async def test_same_compression_still_compresses_routing_when_nothing_armed_it(self, registered_guardrail):
|
||||
"""Regression: only the proxy calls arm_pre_call. Used through the SDK, nothing
|
||||
arms the model-side guardrail and nothing has compressed anything, so reusing a
|
||||
model-hop result that was never produced would serve the request with no
|
||||
compression on either hop, silently ignoring the configuration."""
|
||||
from litellm.proxy.guardrails import auto_router_compression
|
||||
|
||||
router, strategy = self._router(
|
||||
{
|
||||
"auto_router_routing_compression": "fake-compress",
|
||||
"auto_router_model_compression": "fake-compress",
|
||||
}
|
||||
)
|
||||
uncompressed = self._messages()
|
||||
assert auto_router_compression.model_hop_compression_armed() is False
|
||||
|
||||
await router.async_pre_routing_hook(
|
||||
model="smart-router", request_kwargs={"metadata": {}}, messages=uncompressed
|
||||
)
|
||||
|
||||
assert strategy.received_messages != uncompressed
|
||||
assert registered_guardrail.call_count == 1
|
||||
|
||||
async def test_same_compression_on_both_hops_compresses_once(self, registered_guardrail):
|
||||
"""The same/different distinction exists so a shared choice does not pay for
|
||||
compression twice: by the time the router runs, `messages` already reflects
|
||||
whatever the ordinary pre-call guardrail pipeline did for the model call, so
|
||||
the routing decision must reuse it rather than calling the guardrail again."""
|
||||
from litellm.proxy.guardrails import auto_router_compression
|
||||
|
||||
router, strategy = self._router(
|
||||
{
|
||||
"auto_router_routing_compression": "fake-compress",
|
||||
"auto_router_model_compression": "fake-compress",
|
||||
}
|
||||
)
|
||||
# Stands in for what the proxy's ordinary pre-call guardrail pipeline would
|
||||
# have already produced for the model call, since `auto_router_model_compression`
|
||||
# names a guardrail: the router never triggers that pipeline itself.
|
||||
already_compressed_messages = [{"role": "user", "content": "[COMPRESSED] What is the capital of France?"}]
|
||||
# arm_pre_call is what would have armed that guardrail, and only the proxy calls
|
||||
# it; the reuse below is conditional on it having run.
|
||||
armed = auto_router_compression._model_hop_armed.set(True)
|
||||
|
||||
try:
|
||||
response = await router.async_pre_routing_hook(
|
||||
model="smart-router", request_kwargs={"metadata": {}}, messages=already_compressed_messages
|
||||
)
|
||||
finally:
|
||||
auto_router_compression._model_hop_armed.reset(armed)
|
||||
|
||||
assert strategy.received_messages == already_compressed_messages
|
||||
assert response.messages == already_compressed_messages
|
||||
assert registered_guardrail.call_count == 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_no_policy_is_fully_unaffected(self, registered_guardrail):
|
||||
router, strategy = self._router({})
|
||||
original_messages = self._messages()
|
||||
|
||||
response = await router.async_pre_routing_hook(
|
||||
model="smart-router", request_kwargs={"metadata": {}}, messages=original_messages
|
||||
)
|
||||
|
||||
assert strategy.received_messages is original_messages
|
||||
assert response.messages == original_messages
|
||||
assert registered_guardrail.call_count == 0
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("local_model_cost_map")
|
||||
|
||||
@pytest.mark.usefixtures("local_model_cost_map")
|
||||
class TestAzureBaseModelFallbackLogging:
|
||||
"""When an azure deployment has no base_model but its model name is a known
|
||||
|
|
|
|||
|
|
@ -1,11 +1,11 @@
|
|||
import { SimpleTooltip } from "@/components/ui/tooltip";
|
||||
import { MultiSelect } from "@/components/shared/MultiSelect";
|
||||
import { SearchSelect } from "@/components/shared/SearchSelect";
|
||||
import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select";
|
||||
import { ChevronRight, Info, Plus, Trash2, X } from "lucide-react";
|
||||
import { Switch } from "@/components/ui/switch";
|
||||
|
||||
import { AffinityControls } from "./AffinityControls";
|
||||
import TierRowSelect from "./TierRowSelect";
|
||||
import { ModalityRoutingControls } from "./ModalityRoutingControls";
|
||||
import { Card, CardContent } from "@/components/ui/card";
|
||||
import { Collapsible, CollapsibleContent, CollapsibleTrigger } from "@/components/ui/collapsible";
|
||||
|
|
@ -50,6 +50,8 @@ import EscalationKeywords from "./EscalationKeywords";
|
|||
import KeywordTierRules, { KeywordTierRule } from "./KeywordTierRules";
|
||||
import SemanticKeywordMatching from "./SemanticKeywordMatching";
|
||||
import { type DimensionWeights, type TierBoundaries, type TokenThresholds } from "./heuristic_scoring_knobs";
|
||||
import CompressionControls from "./CompressionControls";
|
||||
import { type AutoRouterCompressionState, DEFAULT_AUTO_ROUTER_COMPRESSION } from "./buildAutoRouterCompression";
|
||||
|
||||
export type { DimensionWeights, TierBoundaries, TokenThresholds };
|
||||
export type { CustomTierSet, TierRow } from "./tier_rows";
|
||||
|
|
@ -370,27 +372,6 @@ const TierRowEditFields: React.FC<{
|
|||
</>
|
||||
);
|
||||
|
||||
const TierRowSelect: React.FC<{
|
||||
label: string;
|
||||
options: { value: string; label: string }[];
|
||||
value: string | null;
|
||||
onValueChange: (rowId: string) => void;
|
||||
placeholder?: string;
|
||||
}> = ({ label, options, value, onValueChange, placeholder }) => (
|
||||
<Select items={options} value={value} onValueChange={(rowId: string | null) => rowId && onValueChange(rowId)}>
|
||||
<SelectTrigger aria-label={label} className="w-full">
|
||||
<SelectValue placeholder={placeholder} />
|
||||
</SelectTrigger>
|
||||
<SelectContent>
|
||||
{options.map((option) => (
|
||||
<SelectItem key={option.value} value={option.value}>
|
||||
{option.label}
|
||||
</SelectItem>
|
||||
))}
|
||||
</SelectContent>
|
||||
</Select>
|
||||
);
|
||||
|
||||
export type AdaptiveEligible = "all" | "classified_tier";
|
||||
|
||||
export type ComplexityTierLabels = Partial<Record<keyof ComplexityTiers, string>>;
|
||||
|
|
@ -502,6 +483,10 @@ interface ComplexityRouterConfigProps {
|
|||
onMatchThresholdChange?: (threshold: number) => void;
|
||||
escalationKeywords?: string[];
|
||||
onEscalationKeywordsChange?: (keywords: string[]) => void;
|
||||
// Optional: not part of complexity_router_config, since it applies to every
|
||||
// pre-routing strategy, not just the complexity router.
|
||||
autoRouterCompression?: AutoRouterCompressionState;
|
||||
onAutoRouterCompressionChange?: (state: AutoRouterCompressionState) => void;
|
||||
showValidationErrors?: boolean;
|
||||
}
|
||||
|
||||
|
|
@ -604,6 +589,8 @@ const ComplexityRouterConfig: React.FC<ComplexityRouterConfigProps> = ({
|
|||
onMatchThresholdChange = () => {},
|
||||
escalationKeywords = [],
|
||||
onEscalationKeywordsChange,
|
||||
autoRouterCompression = DEFAULT_AUTO_ROUTER_COMPRESSION,
|
||||
onAutoRouterCompressionChange,
|
||||
showValidationErrors = false,
|
||||
}) => {
|
||||
const customTierSet = value.custom_tier_set;
|
||||
|
|
@ -877,6 +864,17 @@ const ComplexityRouterConfig: React.FC<ComplexityRouterConfigProps> = ({
|
|||
},
|
||||
]
|
||||
: []),
|
||||
...(onAutoRouterCompressionChange
|
||||
? [
|
||||
{
|
||||
key: "compression",
|
||||
label: <strong className="text-foreground font-semibold">Advanced: Compression</strong>,
|
||||
children: (
|
||||
<CompressionControls value={autoRouterCompression} onChange={onAutoRouterCompressionChange} />
|
||||
),
|
||||
},
|
||||
]
|
||||
: []),
|
||||
...(onKeywordTierRulesChange || onSemanticMatchingEnabledChange
|
||||
? [
|
||||
{
|
||||
|
|
|
|||
|
|
@ -0,0 +1,91 @@
|
|||
import { SimpleTooltip } from "@/components/ui/tooltip";
|
||||
import { SearchSelect, SearchSelectOption } from "@/components/shared/SearchSelect";
|
||||
import { Label } from "@/components/ui/label";
|
||||
import { RadioGroup, RadioGroupItem } from "@/components/ui/radio-group";
|
||||
import { Info } from "lucide-react";
|
||||
import React from "react";
|
||||
import { useGuardrails } from "@/app/(dashboard)/hooks/guardrails/useGuardrails";
|
||||
import {
|
||||
AutoRouterCompressionState,
|
||||
isCompressionGuardrailProvider,
|
||||
NO_COMPRESSION,
|
||||
} from "./buildAutoRouterCompression";
|
||||
|
||||
interface CompressionControlsProps {
|
||||
value: AutoRouterCompressionState;
|
||||
onChange: (state: AutoRouterCompressionState) => void;
|
||||
}
|
||||
|
||||
const NONE_OPTION: SearchSelectOption = { label: "None (no compression)", value: NO_COMPRESSION };
|
||||
|
||||
const CompressionControls: React.FC<CompressionControlsProps> = ({ value, onChange }) => {
|
||||
const { routing, sameAsRouting, model } = value;
|
||||
const onRoutingChange = (newRouting: string | undefined) =>
|
||||
onChange({ ...value, routing: newRouting, sameAsRouting: newRouting === undefined ? true : sameAsRouting });
|
||||
const onSameAsRoutingChange = (newSameAsRouting: boolean) => onChange({ ...value, sameAsRouting: newSameAsRouting });
|
||||
const onModelChange = (newModel: string | undefined) => onChange({ ...value, model: newModel });
|
||||
|
||||
const { data } = useGuardrails();
|
||||
const compressionOptions: SearchSelectOption[] = (data?.guardrails ?? [])
|
||||
.filter((g) => isCompressionGuardrailProvider(g.litellm_params?.guardrail))
|
||||
.map((g) => ({ label: g.guardrail_name, value: g.guardrail_name }));
|
||||
const options: SearchSelectOption[] = [NONE_OPTION, ...compressionOptions];
|
||||
|
||||
return (
|
||||
<div className="space-y-4">
|
||||
<div>
|
||||
<div className="mb-1 flex items-center gap-2">
|
||||
<span className="text-sm font-medium">Routing decision</span>
|
||||
<SimpleTooltip content="Compression applied to the classifier's own call that picks a tier, separate from the model the request routes to.">
|
||||
<Info className="size-4 text-muted-foreground" />
|
||||
</SimpleTooltip>
|
||||
</div>
|
||||
<SearchSelect
|
||||
options={options}
|
||||
value={routing ?? ""}
|
||||
onValueChange={(value) => onRoutingChange(value === "" ? undefined : value)}
|
||||
placeholder="Inherit from the request's own compression guardrails"
|
||||
emptyText="No compression guardrails found"
|
||||
aria-label="Routing decision compression"
|
||||
/>
|
||||
</div>
|
||||
|
||||
{routing !== undefined && (
|
||||
<div>
|
||||
<span className="mb-2 block text-sm font-medium">Model call</span>
|
||||
<RadioGroup
|
||||
value={sameAsRouting ? "same" : "different"}
|
||||
onValueChange={(value: unknown) => onSameAsRoutingChange(value === "same")}
|
||||
className="w-full"
|
||||
>
|
||||
<div className="flex w-full flex-col items-start gap-2">
|
||||
<Label className="items-start font-normal leading-normal">
|
||||
<RadioGroupItem value="same" className="mt-0.5" />
|
||||
<span>Same as the routing decision</span>
|
||||
</Label>
|
||||
<Label className="items-start font-normal leading-normal">
|
||||
<RadioGroupItem value="different" className="mt-0.5" />
|
||||
<span>Use a different compression</span>
|
||||
</Label>
|
||||
</div>
|
||||
</RadioGroup>
|
||||
|
||||
{!sameAsRouting && (
|
||||
<div className="mt-3">
|
||||
<SearchSelect
|
||||
options={options}
|
||||
value={model ?? ""}
|
||||
onValueChange={(value) => onModelChange(value === "" ? undefined : value)}
|
||||
placeholder="None (no compression)"
|
||||
emptyText="No compression guardrails found"
|
||||
aria-label="Model call compression"
|
||||
/>
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
);
|
||||
};
|
||||
|
||||
export default CompressionControls;
|
||||
|
|
@ -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 }) => (
|
||||
<Select items={options} value={value} onValueChange={(rowId: string | null) => rowId && onValueChange(rowId)}>
|
||||
<SelectTrigger aria-label={label} className="w-full">
|
||||
<SelectValue placeholder={placeholder} />
|
||||
</SelectTrigger>
|
||||
<SelectContent>
|
||||
{options.map((option) => (
|
||||
<SelectItem key={option.value} value={option.value}>
|
||||
{option.label}
|
||||
</SelectItem>
|
||||
))}
|
||||
</SelectContent>
|
||||
</Select>
|
||||
);
|
||||
|
||||
export default TierRowSelect;
|
||||
|
|
@ -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(<Harness />);
|
||||
|
||||
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(<Harness />);
|
||||
|
||||
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(<Harness />);
|
||||
|
||||
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 () => {
|
||||
|
|
|
|||
|
|
@ -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<AddAutoRouterTabProps> = ({
|
|||
const [embeddingModel, setEmbeddingModel] = useState<string | undefined>(undefined);
|
||||
const [matchThreshold, setMatchThreshold] = useState<number>(DEFAULT_MATCH_THRESHOLD);
|
||||
const [escalationKeywords, setEscalationKeywords] = useState<string[]>(DEFAULT_ESCALATION_KEYWORDS);
|
||||
const [autoRouterCompression, setAutoRouterCompression] = useState<AutoRouterCompressionState>(
|
||||
DEFAULT_AUTO_ROUTER_COMPRESSION,
|
||||
);
|
||||
const [showValidationErrors, setShowValidationErrors] = useState<boolean>(false);
|
||||
const [editingTiers, setEditingTiers] = useState(false);
|
||||
const [isSubmitting, setIsSubmitting] = useState(false);
|
||||
|
|
@ -465,6 +473,7 @@ const AddAutoRouterTab: React.FC<AddAutoRouterTabProps> = ({
|
|||
model_type: "complexity_router",
|
||||
complexity_router_config: complexityRouterConfigPayload,
|
||||
model_access_group: form.getValues("model_access_group"),
|
||||
...buildAutoRouterCompressionParams(autoRouterCompression),
|
||||
};
|
||||
|
||||
await handleAddAutoRouterSubmit(submitValues, accessToken, () => form.reset(EMPTY_FORM_VALUES), handleOk);
|
||||
|
|
@ -670,6 +679,8 @@ const AddAutoRouterTab: React.FC<AddAutoRouterTabProps> = ({
|
|||
onMatchThresholdChange={setMatchThreshold}
|
||||
escalationKeywords={escalationKeywords}
|
||||
onEscalationKeywordsChange={setEscalationKeywords}
|
||||
autoRouterCompression={autoRouterCompression}
|
||||
onAutoRouterCompressionChange={setAutoRouterCompression}
|
||||
showValidationErrors={showValidationErrors}
|
||||
/>
|
||||
</div>
|
||||
|
|
|
|||
|
|
@ -0,0 +1,119 @@
|
|||
import {
|
||||
buildAutoRouterCompressionParams,
|
||||
DEFAULT_AUTO_ROUTER_COMPRESSION,
|
||||
hydrateAutoRouterCompression,
|
||||
NO_COMPRESSION,
|
||||
} from "./buildAutoRouterCompression";
|
||||
|
||||
describe("buildAutoRouterCompressionParams", () => {
|
||||
it("omits both keys when routing was never configured", () => {
|
||||
expect(buildAutoRouterCompressionParams(DEFAULT_AUTO_ROUTER_COMPRESSION)).toEqual({});
|
||||
});
|
||||
|
||||
it("mirrors routing onto model when same-as-routing is chosen", () => {
|
||||
const params = buildAutoRouterCompressionParams({
|
||||
routing: "headroom-a",
|
||||
sameAsRouting: true,
|
||||
model: undefined,
|
||||
});
|
||||
expect(params).toEqual({
|
||||
auto_router_routing_compression: "headroom-a",
|
||||
auto_router_model_compression: "headroom-a",
|
||||
});
|
||||
});
|
||||
|
||||
it("uses the explicit model choice when different is chosen", () => {
|
||||
const params = buildAutoRouterCompressionParams({
|
||||
routing: "headroom-a",
|
||||
sameAsRouting: false,
|
||||
model: "headroom-b",
|
||||
});
|
||||
expect(params).toEqual({
|
||||
auto_router_routing_compression: "headroom-a",
|
||||
auto_router_model_compression: "headroom-b",
|
||||
});
|
||||
});
|
||||
|
||||
it("defaults the model side to none when different is chosen but nothing is picked", () => {
|
||||
const params = buildAutoRouterCompressionParams({
|
||||
routing: "headroom-a",
|
||||
sameAsRouting: false,
|
||||
model: undefined,
|
||||
});
|
||||
expect(params).toEqual({
|
||||
auto_router_routing_compression: "headroom-a",
|
||||
auto_router_model_compression: NO_COMPRESSION,
|
||||
});
|
||||
});
|
||||
|
||||
it("sends the none sentinel when routing itself is explicitly turned off", () => {
|
||||
const params = buildAutoRouterCompressionParams({
|
||||
routing: NO_COMPRESSION,
|
||||
sameAsRouting: true,
|
||||
model: undefined,
|
||||
});
|
||||
expect(params).toEqual({
|
||||
auto_router_routing_compression: NO_COMPRESSION,
|
||||
auto_router_model_compression: NO_COMPRESSION,
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
describe("hydrateAutoRouterCompression", () => {
|
||||
it("returns the default state when neither key is set", () => {
|
||||
expect(hydrateAutoRouterCompression({})).toEqual(DEFAULT_AUTO_ROUTER_COMPRESSION);
|
||||
});
|
||||
|
||||
it("is same-as-routing when the model value matches routing", () => {
|
||||
const state = hydrateAutoRouterCompression({
|
||||
auto_router_routing_compression: "headroom-a",
|
||||
auto_router_model_compression: "headroom-a",
|
||||
});
|
||||
expect(state).toEqual({ routing: "headroom-a", sameAsRouting: true, model: undefined });
|
||||
});
|
||||
|
||||
it("is different when the model value diverges from routing", () => {
|
||||
const state = hydrateAutoRouterCompression({
|
||||
auto_router_routing_compression: "headroom-a",
|
||||
auto_router_model_compression: "headroom-b",
|
||||
});
|
||||
expect(state).toEqual({ routing: "headroom-a", sameAsRouting: false, model: "headroom-b" });
|
||||
});
|
||||
|
||||
it("treats a missing model key as no model-hop compression, not same-as-routing", () => {
|
||||
const state = hydrateAutoRouterCompression({ auto_router_routing_compression: "headroom-a" });
|
||||
expect(state).toEqual({ routing: "headroom-a", sameAsRouting: false, model: "none" });
|
||||
});
|
||||
|
||||
it("re-saving a routing-only config leaves the model hop uncompressed", () => {
|
||||
// Regression: the backend reads an absent model key as no model-hop compression.
|
||||
// Hydrating it as same-as-routing made opening the router and saving any unrelated
|
||||
// edit write the routing guardrail onto the model hop, so the model call silently
|
||||
// started receiving compressed messages.
|
||||
const stored = { auto_router_routing_compression: "headroom-a" };
|
||||
const rebuilt = buildAutoRouterCompressionParams(hydrateAutoRouterCompression(stored));
|
||||
expect(rebuilt.auto_router_model_compression).toBe("none");
|
||||
expect(rebuilt.auto_router_model_compression).not.toBe("headroom-a");
|
||||
});
|
||||
|
||||
it("surfaces a stored model-only policy instead of reading as untouched", () => {
|
||||
// Regression: the backend treats either key alone as an authoritative policy, so a
|
||||
// model-only config that hydrated to the inherit state was invisible in the form,
|
||||
// and the next save overwrote the stored model hop with the routing value.
|
||||
const state = hydrateAutoRouterCompression({ auto_router_model_compression: "headroom-b" });
|
||||
expect(state).toEqual({ routing: "none", sameAsRouting: false, model: "headroom-b" });
|
||||
});
|
||||
|
||||
it("round-trips a model-only policy without changing either hop", () => {
|
||||
const stored = { auto_router_model_compression: "headroom-b" };
|
||||
const rebuilt = buildAutoRouterCompressionParams(hydrateAutoRouterCompression(stored));
|
||||
expect(rebuilt.auto_router_model_compression).toBe("headroom-b");
|
||||
expect(rebuilt.auto_router_routing_compression).toBe("none");
|
||||
});
|
||||
|
||||
it("round-trips through buildAutoRouterCompressionParams", () => {
|
||||
const original = { auto_router_routing_compression: "headroom-a", auto_router_model_compression: "none" };
|
||||
const rebuilt = buildAutoRouterCompressionParams(hydrateAutoRouterCompression(original));
|
||||
expect(rebuilt).toEqual(original);
|
||||
});
|
||||
});
|
||||
|
|
@ -0,0 +1,69 @@
|
|||
/**
|
||||
* Maps the auto router's compression form state to the two flat litellm_params keys
|
||||
* the backend reads (litellm.proxy.guardrails.auto_router_compression), and back.
|
||||
*
|
||||
* `routing` being undefined means the section was never touched: both keys are
|
||||
* omitted from the payload, and the request's own compression guardrails apply to
|
||||
* both hops unchanged. Once `routing` has a value (a guardrail name, or the "none"
|
||||
* sentinel for explicit no-compression), the auto router is authoritative and the
|
||||
* model side always gets a concrete value too, mirroring `routing` when same-as
|
||||
* is chosen and defaulting to "none" otherwise.
|
||||
*/
|
||||
|
||||
export const NO_COMPRESSION = "none";
|
||||
|
||||
/** Guardrail providers that compress prompts, mirroring COMPRESSION_GUARDRAIL_PROVIDERS in
|
||||
* litellm/proxy/guardrails/auto_router_compression.py. Both are selectable per hop. */
|
||||
export const COMPRESSION_GUARDRAIL_PROVIDERS: readonly string[] = ["headroom", "compresr"];
|
||||
|
||||
export const isCompressionGuardrailProvider = (provider: unknown): boolean =>
|
||||
typeof provider === "string" && COMPRESSION_GUARDRAIL_PROVIDERS.includes(provider.toLowerCase());
|
||||
|
||||
export interface AutoRouterCompressionState {
|
||||
routing: string | undefined;
|
||||
sameAsRouting: boolean;
|
||||
model: string | undefined;
|
||||
}
|
||||
|
||||
export interface AutoRouterCompressionLitellmParams {
|
||||
auto_router_routing_compression?: string;
|
||||
auto_router_model_compression?: string;
|
||||
}
|
||||
|
||||
export const DEFAULT_AUTO_ROUTER_COMPRESSION: AutoRouterCompressionState = {
|
||||
routing: undefined,
|
||||
sameAsRouting: true,
|
||||
model: undefined,
|
||||
};
|
||||
|
||||
export const buildAutoRouterCompressionParams = (
|
||||
state: AutoRouterCompressionState,
|
||||
): AutoRouterCompressionLitellmParams => {
|
||||
if (state.routing === undefined) return {};
|
||||
return {
|
||||
auto_router_routing_compression: state.routing,
|
||||
auto_router_model_compression: state.sameAsRouting ? state.routing : state.model ?? NO_COMPRESSION,
|
||||
};
|
||||
};
|
||||
|
||||
export const hydrateAutoRouterCompression = (litellmParams: {
|
||||
auto_router_routing_compression?: string | null;
|
||||
auto_router_model_compression?: string | null;
|
||||
}): AutoRouterCompressionState => {
|
||||
const storedRouting = litellmParams.auto_router_routing_compression ?? undefined;
|
||||
const storedModel = litellmParams.auto_router_model_compression ?? undefined;
|
||||
|
||||
// Only neither key set means the section was never touched. The backend treats
|
||||
// either key on its own as an authoritative policy (policy_from_litellm_params), so
|
||||
// reading a model-only config as untouched would hide it from the form and let the
|
||||
// next save overwrite the stored model hop.
|
||||
if (storedRouting === undefined && storedModel === undefined) return DEFAULT_AUTO_ROUTER_COMPRESSION;
|
||||
|
||||
// An absent key on either hop is no compression for that hop, not same-as-the-other:
|
||||
// the backend reads it as None. Hydrating it as same-as-routing would make re-saving
|
||||
// an unrelated edit write one hop's guardrail onto the other.
|
||||
const routing = storedRouting ?? NO_COMPRESSION;
|
||||
const model = storedModel ?? NO_COMPRESSION;
|
||||
const sameAsRouting = model === routing;
|
||||
return { routing, sameAsRouting, model: sameAsRouting ? undefined : model };
|
||||
};
|
||||
|
|
@ -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 } : {}),
|
||||
|
|
|
|||
|
|
@ -1036,6 +1036,87 @@ describe("EditAutoRouterModal with a stored custom tier set", () => {
|
|||
expect(savedConfig().tier_model_configs).toEqual(CUSTOM_STORED.tier_model_configs);
|
||||
});
|
||||
});
|
||||
describe("EditAutoRouterModal prompt compression", () => {
|
||||
beforeEach(() => {
|
||||
modelPatchUpdateCall.mockClear();
|
||||
});
|
||||
|
||||
const savedLitellmParams = () => {
|
||||
const [, payload] = modelPatchUpdateCall.mock.calls.at(-1) ?? [];
|
||||
return payload?.litellm_params;
|
||||
};
|
||||
|
||||
const renderWithStoredCompression = (compression?: {
|
||||
auto_router_routing_compression?: string;
|
||||
auto_router_model_compression?: string;
|
||||
}) =>
|
||||
renderWithProviders(
|
||||
<EditAutoRouterModal
|
||||
isVisible
|
||||
onCancel={vi.fn()}
|
||||
onSuccess={vi.fn()}
|
||||
modelData={{
|
||||
...MODEL_DATA,
|
||||
litellm_params: { ...MODEL_DATA.litellm_params, ...compression },
|
||||
}}
|
||||
accessToken="token"
|
||||
userRole="Admin"
|
||||
/>,
|
||||
);
|
||||
|
||||
it("leaves both compression keys out of an untouched save when none were stored", async () => {
|
||||
const user = userEvent.setup();
|
||||
renderWithStoredCompression();
|
||||
|
||||
await user.click(await screen.findByRole("button", { name: /save changes/i }));
|
||||
|
||||
await waitFor(() => expect(modelPatchUpdateCall).toHaveBeenCalled());
|
||||
expect(savedLitellmParams()).not.toHaveProperty("auto_router_routing_compression");
|
||||
expect(savedLitellmParams()).not.toHaveProperty("auto_router_model_compression");
|
||||
});
|
||||
|
||||
it("preserves a stored same-as-routing compression through an untouched open-and-save", async () => {
|
||||
const user = userEvent.setup();
|
||||
renderWithStoredCompression({
|
||||
auto_router_routing_compression: "headroom-a",
|
||||
auto_router_model_compression: "headroom-a",
|
||||
});
|
||||
|
||||
await user.click(await screen.findByRole("button", { name: /save changes/i }));
|
||||
|
||||
await waitFor(() => expect(modelPatchUpdateCall).toHaveBeenCalled());
|
||||
expect(savedLitellmParams()?.auto_router_routing_compression).toBe("headroom-a");
|
||||
expect(savedLitellmParams()?.auto_router_model_compression).toBe("headroom-a");
|
||||
});
|
||||
|
||||
it("shows a stored different-compression choice as Use a different compression, not Same", async () => {
|
||||
const user = userEvent.setup();
|
||||
renderWithStoredCompression({
|
||||
auto_router_routing_compression: "headroom-a",
|
||||
auto_router_model_compression: "none",
|
||||
});
|
||||
|
||||
await user.click(await screen.findByText("Advanced: Compression"));
|
||||
|
||||
expect(await screen.findByRole("combobox", { name: "Routing decision compression" })).toHaveValue("headroom-a");
|
||||
expect(screen.getByRole("radio", { name: "Use a different compression" })).toBeChecked();
|
||||
expect(screen.getByRole("combobox", { name: "Model call compression" })).toHaveValue("None (no compression)");
|
||||
});
|
||||
|
||||
it("preserves a stored different-compression choice through an untouched open-and-save", async () => {
|
||||
const user = userEvent.setup();
|
||||
renderWithStoredCompression({
|
||||
auto_router_routing_compression: "headroom-a",
|
||||
auto_router_model_compression: "none",
|
||||
});
|
||||
|
||||
await user.click(await screen.findByRole("button", { name: /save changes/i }));
|
||||
|
||||
await waitFor(() => expect(modelPatchUpdateCall).toHaveBeenCalled());
|
||||
expect(savedLitellmParams()?.auto_router_routing_compression).toBe("headroom-a");
|
||||
expect(savedLitellmParams()?.auto_router_model_compression).toBe("none");
|
||||
});
|
||||
});
|
||||
|
||||
describe("EditAutoRouterModal classifier vision", () => {
|
||||
beforeEach(() => {
|
||||
|
|
|
|||
|
|
@ -41,6 +41,12 @@ import {
|
|||
} from "../add_model/build_complexity_router_config";
|
||||
import { KeywordTierRule } from "../add_model/KeywordTierRules";
|
||||
import { DEFAULT_MATCH_THRESHOLD } from "../add_model/SemanticKeywordMatching";
|
||||
import {
|
||||
type AutoRouterCompressionState,
|
||||
buildAutoRouterCompressionParams,
|
||||
DEFAULT_AUTO_ROUTER_COMPRESSION,
|
||||
hydrateAutoRouterCompression,
|
||||
} from "../add_model/buildAutoRouterCompression";
|
||||
import { hydrateKeywordTierRules } from "../add_model/complexity_router_keywords";
|
||||
import {
|
||||
hydrateDimensionWeights,
|
||||
|
|
@ -447,6 +453,9 @@ const EditAutoRouterModal: React.FC<EditAutoRouterModalProps> = ({
|
|||
const [semanticMatchingEnabled, setSemanticMatchingEnabled] = useState<boolean>(false);
|
||||
const [embeddingModel, setEmbeddingModel] = useState<string | undefined>(undefined);
|
||||
const [matchThreshold, setMatchThreshold] = useState<number>(DEFAULT_MATCH_THRESHOLD);
|
||||
const [autoRouterCompression, setAutoRouterCompression] = useState<AutoRouterCompressionState>(
|
||||
DEFAULT_AUTO_ROUTER_COMPRESSION,
|
||||
);
|
||||
const [complexityRouterConfig, setComplexityRouterConfig] = useState<ComplexityRouterConfigValue>({
|
||||
tiers: { SIMPLE: [], MEDIUM: [], COMPLEX: [], REASONING: [] },
|
||||
classifier_type: "heuristic",
|
||||
|
|
@ -539,6 +548,12 @@ const EditAutoRouterModal: React.FC<EditAutoRouterModalProps> = ({
|
|||
setMatchThreshold(
|
||||
typeof parsedConfig.match_threshold === "number" ? parsedConfig.match_threshold : DEFAULT_MATCH_THRESHOLD,
|
||||
);
|
||||
setAutoRouterCompression(
|
||||
hydrateAutoRouterCompression({
|
||||
auto_router_routing_compression: modelData.litellm_params?.auto_router_routing_compression,
|
||||
auto_router_model_compression: modelData.litellm_params?.auto_router_model_compression,
|
||||
}),
|
||||
);
|
||||
|
||||
form.reset({
|
||||
...EMPTY_FORM_VALUES,
|
||||
|
|
@ -651,6 +666,7 @@ const EditAutoRouterModal: React.FC<EditAutoRouterModalProps> = ({
|
|||
...modelData.litellm_params,
|
||||
complexity_router_config: updatedConfig,
|
||||
complexity_router_default_model: defaultModel,
|
||||
...buildAutoRouterCompressionParams(autoRouterCompression),
|
||||
};
|
||||
const updatedModelInfo = {
|
||||
...modelData.model_info,
|
||||
|
|
@ -772,6 +788,8 @@ const EditAutoRouterModal: React.FC<EditAutoRouterModalProps> = ({
|
|||
onMatchThresholdChange={setMatchThreshold}
|
||||
escalationKeywords={escalationKeywords}
|
||||
onEscalationKeywordsChange={setEscalationKeywords}
|
||||
autoRouterCompression={autoRouterCompression}
|
||||
onAutoRouterCompressionChange={setAutoRouterCompression}
|
||||
/>
|
||||
</div>
|
||||
) : (
|
||||
|
|
|
|||
8
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
8
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -29281,6 +29281,10 @@ export interface components {
|
|||
auto_router_embedding_model?: string | null;
|
||||
/** Auto Router Max Input Chars */
|
||||
auto_router_max_input_chars?: number | null;
|
||||
/** Auto Router Model Compression */
|
||||
auto_router_model_compression?: string | null;
|
||||
/** Auto Router Routing Compression */
|
||||
auto_router_routing_compression?: string | null;
|
||||
/** Aws Access Key Id */
|
||||
aws_access_key_id?: string | null;
|
||||
/** Aws Batch Role Arn */
|
||||
|
|
@ -39411,6 +39415,10 @@ export interface components {
|
|||
auto_router_embedding_model?: string | null;
|
||||
/** Auto Router Max Input Chars */
|
||||
auto_router_max_input_chars?: number | null;
|
||||
/** Auto Router Model Compression */
|
||||
auto_router_model_compression?: string | null;
|
||||
/** Auto Router Routing Compression */
|
||||
auto_router_routing_compression?: string | null;
|
||||
/** Aws Access Key Id */
|
||||
aws_access_key_id?: string | null;
|
||||
/** Aws Batch Role Arn */
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue