scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-13 08:52:02 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 353

feat(proxy): let pre_call_hook skip built-in guardrails for a model routing refuses

Browse source
This commit is contained in:
yucheng 2026-10-30 11:31:19 -07:00
parent 2e2a46aed3
commit 4cd101a61c
2 changed files with 292 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

103
litellm/proxy/utils.py
View file

@ -311,6 +311,7 @@ if TYPE_CHECKING:
from litellm.proxy.db.baseline_accounting import BaselineAccountingRecord
from litellm.proxy.db.model_usage_rollup import ModelUsageTransaction
from litellm.proxy.db.spend_log_tool_index import ToolUsageTransaction
from litellm.proxy.route_llm_request import UnroutableModelProbe
from litellm.repositories.prisma_protocols import TableActions
from litellm.types.mcp_server.mcp_server_manager import MCPServer
from litellm.types.proxy.policy_engine.pipeline_types import GuardrailPipeline, PipelineStep
@ -1659,6 +1660,8 @@ class _CallbackCapabilities:
# every string entry in ``litellm.callbacks``.
resolved_callbacks: tuple[object, ...] = field(default_factory=tuple)
listed_models_filters: tuple[CustomLogger, ...] = field(default_factory=tuple)
last_custom_pre_call_hook: int = -1
routes_on_sensitive_data: bool = False
def _overrides_hook(callback: CustomLogger, hook_name: str) -> bool:
@ -1670,6 +1673,63 @@ def _overrides_moderation_hook(callback: CustomLogger) -> bool:
return _overrides_hook(callback, "async_moderation_hook")
_LITELLM_PACKAGES: Final = frozenset({"litellm", "litellm_enterprise", "enterprise"})
def _is_custom_pre_call_hook(callback: CustomLogger) -> bool:
return type(callback).__module__.partition(".")[0] not in _LITELLM_PACKAGES and (
isinstance(callback, CustomGuardrail) or "async_pre_call_hook" in type(callback).__dict__
)
async def _skipped_for_unroutable_model(
guardrail: CustomGuardrail,
data: Mapping[str, object],
probe: "UnroutableModelProbe | None",
) -> bool:
return (
probe is not None
and not _is_custom_pre_call_hook(guardrail)
and guardrail.should_run_guardrail(data=data, event_type=GuardrailEventHooks.pre_call) # pyright: ignore[reportUnknownMemberType] # takes the untyped request dict
and await probe.refuses(data)
)
_SYSTEM_INSTRUCTION_KEYS: Final = ("systemInstruction", "system_instruction")
def without_unscanned_content(request: Mapping[str, object]) -> Mapping[str, object]:
"""Every request field a pre-call guardrail can rewrite, each replaced with a redaction marker."""
system_instruction: Final = {"parts": [{"text": REDACTED_BY_LITELLM}]}
redactions: Final[Mapping[str, object]] = {
"messages": [{"role": "user", "content": REDACTED_BY_LITELLM}],
"contents": [{"role": "user", "parts": [{"text": REDACTED_BY_LITELLM}]}],
"prompt": "",
"input": "",
"query": "",
"instruction": "",
"tools": [],
"system": REDACTED_BY_LITELLM,
"instructions": REDACTED_BY_LITELLM,
**{key: system_instruction for key in _SYSTEM_INSTRUCTION_KEYS},
}
config: Final = request.get("config")
config_fields: Final = _as_str_mapping(config)
return {
**{key: value for key, value in redactions.items() if key in request},
**(
{}
if not isinstance(config, Mapping)
else {
"config": {
**config_fields,
**{key: system_instruction for key in _SYSTEM_INSTRUCTION_KEYS if key in config_fields},
}
}
),
}
_LISTED_MODEL_NAMES: Final = TypeAdapter(tuple[str, ...])
_MCP_TOOL_DESCRIPTION: Final[TypeAdapter[str | None]] = TypeAdapter(str | None)
_MCP_TOOL_INPUT_SCHEMA: Final[TypeAdapter[Mapping[str, object] | None]] = TypeAdapter(Mapping[str, object] | None)
@ -2911,6 +2971,7 @@ class ProxyLogging:
guardrails_only: bool = False,
skip_guardrails: bool = False,
endpoint_type: EndpointType = EndpointType.GENERIC,
routing_probe: "UnroutableModelProbe | None" = None,
) -> None:
pass
@ -2923,6 +2984,7 @@ class ProxyLogging:
guardrails_only: bool = False,
skip_guardrails: bool = False,
endpoint_type: EndpointType = EndpointType.GENERIC,
routing_probe: "UnroutableModelProbe | None" = None,
) -> dict:
pass
@ -2934,6 +2996,7 @@ class ProxyLogging:
guardrails_only: bool = False,
skip_guardrails: bool = False,
endpoint_type: EndpointType = EndpointType.GENERIC,
routing_probe: "UnroutableModelProbe | None" = None,
) -> dict | None:
"""
Allows users to modify/reject the incoming request to the proxy, without having to deal with parsing Request body.
@ -2946,6 +3009,11 @@ class ProxyLogging:
With ``guardrails_only`` the walk is limited to guardrails and guardrail pipelines: rate
limiting, budget accounting, prompt templates and hanging-request alerting are skipped.
Use it to scan a payload that is not itself a request, such as one record of a batch file.
With ``routing_probe``, a pre-call guardrail litellm ships is not called for a model ``route_request``
refuses, when only litellm's own guardrails and hooks come after it and no guardrail reroutes on
sensitive data. The hooks after a skipped sequential guardrail see the request with every field it
could have rewritten redacted.
"""
verbose_proxy_logger.debug("Inside Proxy Logging Pre-call hook!")
@ -3013,6 +3081,7 @@ class ProxyLogging:
)
caps: Final = ProxyLogging.callback_capabilities()
probe: Final = None if caps.routes_on_sensitive_data else routing_probe
# Skip the per-request callback walk entirely when nothing in
# ``litellm.callbacks`` overrides ``async_pre_call_hook`` and no
# CustomGuardrail is configured. Saves the loop overhead +
@ -3040,7 +3109,7 @@ class ProxyLogging:
)
deferred_route_exc: SensitiveDataRouteException | None = None
for _callback in caps.resolved_callbacks:
for position, _callback in enumerate(caps.resolved_callbacks):
start_time = time.time()
try:
if isinstance(_callback, CustomGuardrail) and data is not None:
@ -3054,6 +3123,14 @@ class ProxyLogging:
if getattr(_callback, "run_in_parallel", False):
continue
if position > caps.last_custom_pre_call_hook and await _skipped_for_unroutable_model(
_callback,
data, # pyright: ignore[reportUnknownArgumentType] # data is the untyped request dict
probe,
):
data.update(without_unscanned_content(data)) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # data is the untyped request dict
continue
data = await self._run_sequential_guardrail_callback(
callback=_callback,
data=data,
@ -3120,7 +3197,17 @@ class ProxyLogging:
if parallel_guardrails and data is not None:
await self._run_parallel_pre_call_guardrails(
guardrails=parallel_guardrails,
guardrails=tuple(
[
guardrail
for guardrail in parallel_guardrails
if not await _skipped_for_unroutable_model(
guardrail,
data, # pyright: ignore[reportUnknownArgumentType] # data is the untyped request dict
probe,
)
]
),
data=data,
raw_request_snapshot=raw_request_snapshot,
user_api_key_dict=user_api_key_dict,
@ -3444,6 +3531,18 @@ class ProxyLogging:
listed_models_filters=tuple(
callback for callback in resolved_callbacks if _overrides_hook(callback, "async_filter_listed_models")
),
last_custom_pre_call_hook=max(
(
position
for position, callback in enumerate(resolved_callbacks)
if _is_custom_pre_call_hook(callback)
),
default=-1,
),
routes_on_sensitive_data=any(
isinstance(callback, CustomGuardrail) and callback.should_route_on_sensitive_data()
for callback in resolved_callbacks
),
)
# Limit cache to handle test churn without leaking; production
# callback lists are stable so this rarely grows past 1 entry.

191
tests/unit/proxy/test_proxy_utils.py
View file

@ -9,6 +9,7 @@ from fastapi import HTTPException, Request
from starlette.datastructures import State
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.integrations.custom_logger import CustomLogger
from litellm.proxy.utils import get_docs_url, get_openapi_url, get_redoc_url
from litellm.types.guardrails import GuardrailEventHooks
@ -3716,3 +3717,193 @@ async def test_evict_config_param_clears_the_local_layer_and_survives_a_redis_ou
await evict_config_param("router_settings", cache=cache)
assert await cache.in_memory_cache.async_get_cache(_config_cache_key("router_settings")) is None
class _ShippedGuardrail(CustomGuardrail):
"""Stands in for a guardrail litellm ships: records the messages each call was sent."""
__module__ = "litellm.proxy.guardrails.guardrail_hooks.shipped_double"
def __init__(self, name: str = "shipped", run_in_parallel: bool = False, default_on: bool = True) -> None:
super().__init__(
guardrail_name=name,
event_hook=GuardrailEventHooks.pre_call,
default_on=default_on,
run_in_parallel=run_in_parallel,
)
self.seen: Final[list[object]] = []
async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type):
self.seen.append(data.get("messages"))
return data
class _OperatorGuardrail(_ShippedGuardrail):
"""A guardrail from the operator's own module."""
__module__ = "operator_guardrails.checks"
class _ShippedHook(CustomLogger):
"""Stands in for a pre-call hook litellm ships that is not a guardrail, such as a rate limiter."""
__module__ = "litellm.proxy.hooks.shipped_hook_double"
def __init__(self) -> None:
super().__init__()
self.seen: Final[list[object]] = []
async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type):
self.seen.append(data.get("messages"))
return data
_SECRET_PROMPT: Final = [{"role": "user", "content": "my card is 4111 1111 1111 1111"}]
def _unroutable_model_probe():
from litellm.proxy.route_llm_request import UnroutableModelProbe
return UnroutableModelProbe(
llm_router=Router(
model_list=[
{
"model_name": "gpt-4o-mini",
"litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-test", "mock_response": "pong"},
}
]
),
user_model=None,
route_type="acompletion",
user_api_key_dict=None,
)
async def _walk(monkeypatch, callbacks, model: str, probe) -> dict:
from litellm.caching.caching import DualCache
from litellm.proxy.utils import ProxyLogging
monkeypatch.setattr(litellm, "callbacks", list(callbacks))
ProxyLogging._callback_capabilities_cache.clear()
return await ProxyLogging(user_api_key_cache=DualCache()).pre_call_hook(
user_api_key_dict=UserAPIKeyAuth(api_key="test_key", user_id="test_user"),
data={"model": model, "messages": [dict(message) for message in _SECRET_PROMPT]},
call_type="completion",
routing_probe=probe,
)
@pytest.mark.parametrize("run_in_parallel", [False, True], ids=["sequential", "parallel"])
async def test_pre_call_hook_skips_a_shipped_guardrail_only_for_a_model_routing_refuses(monkeypatch, run_in_parallel):
from litellm.proxy.route_llm_request import ProxyModelNotFoundError
unknown_guardrail: Final = _ShippedGuardrail(run_in_parallel=run_in_parallel)
unknown_probe: Final = _unroutable_model_probe()
await _walk(monkeypatch, [unknown_guardrail], "no-such-model", unknown_probe)
known_guardrail: Final = _ShippedGuardrail(run_in_parallel=run_in_parallel)
known_probe: Final = _unroutable_model_probe()
await _walk(monkeypatch, [known_guardrail], "gpt-4o-mini", known_probe)
assert unknown_guardrail.seen == []
assert isinstance(unknown_probe.refusal, ProxyModelNotFoundError)
assert known_guardrail.seen == [_SECRET_PROMPT]
assert known_probe.refusal is None
async def test_pre_call_hook_without_a_probe_runs_the_shipped_guardrail_for_an_unknown_model(monkeypatch):
guardrail: Final = _ShippedGuardrail()
result: Final = await _walk(monkeypatch, [guardrail], "no-such-model", None)
assert guardrail.seen == [_SECRET_PROMPT]
assert result["messages"] == _SECRET_PROMPT
@pytest.mark.parametrize("operator_in_parallel", [False, True], ids=["sequential", "parallel"])
async def test_pre_call_hook_runs_an_operator_guardrail_and_every_guardrail_before_it(
monkeypatch, operator_in_parallel
):
before: Final = _ShippedGuardrail(name="before")
operator: Final = _OperatorGuardrail(name="operator", run_in_parallel=operator_in_parallel)
after: Final = _ShippedGuardrail(name="after")
await _walk(monkeypatch, [before, operator, after], "no-such-model", _unroutable_model_probe())
assert (before.seen, len(operator.seen), after.seen) == ([_SECRET_PROMPT], 1, [])
async def test_pre_call_hook_runs_every_guardrail_when_one_routes_on_sensitive_data(monkeypatch):
guardrail: Final = _ShippedGuardrail()
sensitive_router: Final = CustomGuardrail(
guardrail_name="sensitive-router",
event_hook=GuardrailEventHooks.pre_call,
on_sensitive_data="route",
sensitive_data_route_to_model="gpt-4o-mini",
)
await _walk(monkeypatch, [guardrail, sensitive_router], "no-such-model", _unroutable_model_probe())
assert guardrail.seen == [_SECRET_PROMPT]
async def test_pre_call_hook_hides_the_prompt_from_hooks_after_a_skipped_guardrail(monkeypatch):
from litellm.constants import REDACTED_BY_LITELLM
hook: Final = _ShippedHook()
result: Final = await _walk(monkeypatch, [_ShippedGuardrail(), hook], "no-such-model", _unroutable_model_probe())
redacted: Final = [{"role": "user", "content": REDACTED_BY_LITELLM}]
assert hook.seen == [redacted]
assert result["messages"] == redacted
async def test_pre_call_hook_leaves_the_request_alone_when_no_guardrail_applies(monkeypatch):
hook: Final = _ShippedHook()
probe: Final = _unroutable_model_probe()
result: Final = await _walk(monkeypatch, [_ShippedGuardrail(default_on=False), hook], "no-such-model", probe)
assert hook.seen == [_SECRET_PROMPT]
assert result["messages"] == _SECRET_PROMPT
assert probe.refusal is None
def test_without_unscanned_content_redacts_every_field_a_guardrail_can_rewrite():
from litellm.constants import REDACTED_BY_LITELLM
from litellm.proxy.utils import without_unscanned_content
instruction: Final = {"parts": [{"text": "card 4111"}]}
redacted_instruction: Final = {"parts": [{"text": REDACTED_BY_LITELLM}]}
assert without_unscanned_content(
{
"model": "gpt-4o-mini",
"temperature": 0,
"messages": _SECRET_PROMPT,
"contents": [{"role": "user", "parts": [{"text": "card 4111"}]}],
"prompt": "card 4111",
"input": ["card 4111"],
"query": "card 4111",
"instruction": "card 4111",
"tools": [{"type": "function", "function": {"name": "card", "description": "4111"}}],
"system": "card 4111",
"instructions": "card 4111",
"systemInstruction": instruction,
"system_instruction": instruction,
"config": {"system_instruction": instruction, "temperature": 0},
}
) == {
"messages": [{"role": "user", "content": REDACTED_BY_LITELLM}],
"contents": [{"role": "user", "parts": [{"text": REDACTED_BY_LITELLM}]}],
"prompt": "",
"input": "",
"query": "",
"instruction": "",
"tools": [],
"system": REDACTED_BY_LITELLM,
"instructions": REDACTED_BY_LITELLM,
"systemInstruction": redacted_instruction,
"system_instruction": redacted_instruction,
"config": {"system_instruction": redacted_instruction, "temperature": 0},
}