scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-14 23:21:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 102

fix(guardrails): log mask when a guardrail adds request keys (#40882)

Browse source 
* fix(guardrails): log mask when a guardrail adds request keys

_inputs_were_modified only compared keys present in the pre-hook baseline, so a
guardrail that injected a new key such as tools was logged as allow. Compare over
the union of both key sets, and narrow the pre_call return value to the same
prompt-bearing keys the baseline holds so passthrough stays allow.

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(guardrails): snapshot apply_guardrail inputs before the hook mutates them

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

---------

Co-authored-by: yucheng <yucheng@berri.ai>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot] 2026-09-12 14:39:40 -07:00 • committed by GitHub
parent 883b722fd2
commit 5b36de4646
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 79 additions and 10 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

29
litellm/integrations/custom_guardrail.py
View file

@ -5,6 +5,7 @@ import os
import secrets
from collections.abc import Mapping
from datetime import datetime
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional, get_args
from litellm._logging import verbose_logger
@ -1279,6 +1280,7 @@ class CustomGuardrail(CustomLogger):
guardrail_response: Final = self._summarize_guardrail_response(
response=response,
original_inputs=original_inputs,
event_type=event_type,
)
verbose_logger.debug("Guardrail response: %s", response)
@ -1298,6 +1300,7 @@ class CustomGuardrail(CustomLogger):
self,
response: object,
original_inputs: Mapping[str, object] | None,
event_type: GuardrailEventHooks | None,
) -> object:
"""Reduce a hook's return value to what is safe to log as ``guardrail_response``.
@ -1305,15 +1308,21 @@ class CustomGuardrail(CustomLogger):
returns the (possibly modified) request payload. Neither is a provider verdict, and
logging them verbatim ships the user's prompt to every logging sink (OTEL spans,
Datadog, spend logs), so both collapse to ``"allow"`` / ``"mask"`` by comparing
against ``original_inputs``, a copy taken before the hook ran. A string result is the
hook's own rejection message (the proxy turns it into a 400), not user input, so it is
logged as is.
against ``original_inputs``, a copy taken before the hook ran. A pre_call baseline only
holds the prompt-bearing keys, so the returned request is narrowed to those same keys
before the comparison. A string result is the hook's own rejection message (the proxy
turns it into a 400), not user input, so it is logged as is.
"""
if response is None:
return {}
if original_inputs is None or not isinstance(response, Mapping):
return response
return "mask" if self._inputs_were_modified(original_inputs, response) else "allow"
compared_response: Final[Mapping[str, object]] = (
MappingProxyType({key: value for key, value in response.items() if key in _PRE_CALL_CONTENT_KEYS})
if event_type == GuardrailEventHooks.pre_call
else response
)
return "mask" if self._inputs_were_modified(original_inputs, compared_response) else "allow"
@staticmethod
def _is_guardrail_intervention(e: Exception) -> bool:
@ -1355,8 +1364,8 @@ class CustomGuardrail(CustomLogger):
raise e
def _inputs_were_modified(self, original_inputs: Mapping[str, object], response: Mapping[str, object]) -> bool:
"""True when any baseline key's value differs in ``response`` (mask), False otherwise (allow)."""
return any(response.get(key) != value for key, value in original_inputs.items())
"""True when any key of either mapping differs between them (mask), False otherwise (allow)."""
return any(original_inputs.get(key) != response.get(key) for key in original_inputs.keys() | response.keys())
def mask_content_in_string(
self,
@ -1476,13 +1485,13 @@ def _original_inputs_for(
) -> dict | None: # mutable-ok: matches _process_response(original_inputs=) signature
"""Baseline the hook's return value is compared against to decide "allow" vs "mask".
``apply_guardrail`` masks a fresh ``inputs`` dict, so that dict is the baseline. Pre-call
hooks edit the request in place and return it, so the baseline is a deep copy of the
prompt-bearing keys taken before the hook runs.
Hooks may edit their argument in place and return it, so the baseline is always a deep
copy taken before the hook runs: the whole ``inputs`` dict for ``apply_guardrail``, the
prompt-bearing request keys for pre-call hooks.
"""
if func_name == "apply_guardrail":
inputs: Final = kwargs.get("inputs")
return inputs if isinstance(inputs, dict) else None
return copy.deepcopy(inputs) if isinstance(inputs, dict) else None
if event_type != GuardrailEventHooks.pre_call:
return None
return {key: copy.deepcopy(value) for key, value in request_data.items() if key in _PRE_CALL_CONTENT_KEYS}

1
litellm/proxy/guardrails/guardrail_hooks/azure/prompt_shield.py
View file

@ -341,6 +341,7 @@ class AzureContentSafetyPromptShieldGuardrail(AzureGuardrailBase, CustomGuardrai
guardrail_response: Final = self._summarize_guardrail_response(
response=response,
original_inputs=original_inputs,
event_type=event_type,
)
self.add_standard_logging_guardrail_information_to_request_data(
guardrail_json_response=guardrail_response,

59
tests/test_litellm/integrations/test_custom_guardrail.py
View file

@ -3016,3 +3016,62 @@ class TestPreCallHookResponseIsNotLoggedVerbatim:
)
assert self._logged_response(data) == "mask"
@pytest.mark.asyncio
async def test_pre_call_hook_adding_tools_logs_mask(self):
class ToolInjectingGuardrail(CustomGuardrail):
@log_guardrail_information
async def async_pre_call_hook(
self,
user_api_key_dict: UserAPIKeyAuth,
cache: object,
data: dict[str, object],
call_type: str,
) -> dict[str, object]:
return {**data, "tools": [{"type": "function", "function": {"name": "guardrail_injected_tool"}}]}
data = self._request()
await ToolInjectingGuardrail(guardrail_name="g").async_pre_call_hook(
user_api_key_dict=UserAPIKeyAuth(), cache=None, data=data, call_type="acompletion"
)
assert self._logged_response(data) == "mask"
@pytest.mark.asyncio
async def test_apply_guardrail_adding_tools_logs_mask(self):
class ToolInjectingGuardrail(CustomGuardrail):
async def apply_guardrail(
self,
inputs: GenericGuardrailAPIInputs,
request_data: dict[str, object],
input_type: Literal["request", "response"],
logging_obj: Optional["LiteLLMLoggingObj"] = None,
) -> GenericGuardrailAPIInputs:
return {**inputs, "tools": [{"type": "function", "function": {"name": "guardrail_injected_tool"}}]}
data = self._request()
await ToolInjectingGuardrail(guardrail_name="g").apply_guardrail(
inputs={"texts": ["SECRET_PROMPT"]}, request_data=data, input_type="request"
)
assert self._logged_response(data) == "mask"
@pytest.mark.asyncio
async def test_apply_guardrail_masking_inputs_in_place_logs_mask(self):
class InPlaceMaskingGuardrail(CustomGuardrail):
async def apply_guardrail(
self,
inputs: GenericGuardrailAPIInputs,
request_data: dict[str, object],
input_type: Literal["request", "response"],
logging_obj: Optional["LiteLLMLoggingObj"] = None,
) -> GenericGuardrailAPIInputs:
inputs["texts"] = ["<REDACTED>"]
return inputs
data = self._request()
await InPlaceMaskingGuardrail(guardrail_name="g").apply_guardrail(
inputs={"texts": ["SECRET_PROMPT"]}, request_data=data, input_type="request"
)
assert self._logged_response(data) == "mask"