scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-15 09:12:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 311

fix(guardrails): snapshot pre_call request before the hook so in-place edits log as mask

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
shivam 2026-09-04 03:01:34 +00:00
parent 24962ef054
commit 49b49f5c52
3 changed files with 85 additions and 46 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

63
litellm/integrations/custom_guardrail.py
View file

@ -1237,8 +1237,6 @@ class CustomGuardrail(CustomLogger):
"""
guardrail_response: Final = self._summarize_guardrail_response(
response=response,
request_data=request_data,
event_type=event_type,
original_inputs=original_inputs,
)
@ -1258,23 +1256,20 @@ class CustomGuardrail(CustomLogger):
def _summarize_guardrail_response(
self,
response: dict | None, # mutable-ok: matches _process_response signature
request_data: dict, # mutable-ok: matches _process_response signature
event_type: GuardrailEventHooks | None,
original_inputs: dict | None, # mutable-ok: matches _process_response signature
original_inputs: Mapping[str, object] | None,
) -> dict | str: # mutable-ok: returns the caller's response unchanged
"""Reduce a hook's return value to what is safe to log as ``guardrail_response``.
``apply_guardrail`` returns the (possibly masked) inputs and ``async_pre_call_hook``
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"``.
Datadog, spend logs), so both collapse to ``"allow"`` / ``"mask"`` by comparing
against ``original_inputs``, a copy taken before the hook ran.
"""
if response is None:
return {}
if original_inputs is not None:
return "mask" if self._inputs_were_modified(original_inputs, response) else "allow"
if event_type == GuardrailEventHooks.pre_call:
return "mask" if self._inputs_were_modified(request_data, response) else "allow"
return response
@staticmethod
@ -1316,24 +1311,9 @@ class CustomGuardrail(CustomLogger):
)
raise e
def _inputs_were_modified(self, original_inputs: dict, response: dict) -> bool:
"""
Compare original inputs with response to determine if content was modified.
Returns True if the inputs were modified (mask scenario), False otherwise (allow scenario).
"""
# Get all keys from both dictionaries
all_keys: Final = set(original_inputs.keys()) | set(response.keys())
# Compare each key's value
for key in all_keys:
original_value = original_inputs.get(key)
response_value = response.get(key)
if original_value != response_value:
return True
# No modifications detected
return False
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())
def mask_content_in_string(
self,
@ -1440,6 +1420,29 @@ def _sync_guardrail_info_to_logging_obj(request_data: dict, logging_obj: object)
_append_slg_to_litellm_params(mcd.get("litellm_params"), entries)
_PRE_CALL_CONTENT_KEYS: Final = frozenset({"messages", "input", "prompt", "system", "instructions", "tools"})
def _original_inputs_for(
func_name: str,
kwargs: Mapping[str, object],
request_data: Mapping[str, object],
event_type: GuardrailEventHooks | None,
) -> 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.
"""
if func_name == "apply_guardrail":
inputs: Final = kwargs.get("inputs")
return 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}
def log_guardrail_information(func):
"""
Decorator to add standard logging guardrail information to any function
@ -1498,9 +1501,7 @@ def log_guardrail_information(func):
event_type: Final = _infer_event_type_from_function_name(func.__name__)
# Store original inputs for comparison (for apply_guardrail functions)
original_inputs = None
if func.__name__ == "apply_guardrail" and "inputs" in kwargs:
original_inputs = kwargs.get("inputs")
original_inputs: Final = _original_inputs_for(func.__name__, kwargs, request_data, event_type)
logging_obj: Final = kwargs.get("logging_obj") or request_data.get("litellm_logging_obj")
self_recorded_token: Final = _guardrail_self_recorded.set(False)
@ -1540,9 +1541,7 @@ def log_guardrail_information(func):
event_type: Final = _infer_event_type_from_function_name(func.__name__)
# Store original inputs for comparison (for apply_guardrail functions)
original_inputs = None
if func.__name__ == "apply_guardrail" and "inputs" in kwargs:
original_inputs = kwargs.get("inputs")
original_inputs: Final = _original_inputs_for(func.__name__, kwargs, request_data, event_type)
logging_obj: Final = kwargs.get("logging_obj") or request_data.get("litellm_logging_obj")
self_recorded_token: Final = _guardrail_self_recorded.set(False)

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

@ -340,8 +340,6 @@ class AzureContentSafetyPromptShieldGuardrail(AzureGuardrailBase, CustomGuardrai
(openai/moderations.py)."""
guardrail_response: Final = self._summarize_guardrail_response(
response=response,
request_data=request_data,
event_type=event_type,
original_inputs=original_inputs,
)
self.add_standard_logging_guardrail_information_to_request_data(

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

@ -2513,22 +2513,37 @@ class TestPreCallHookResponseIsNotLoggedVerbatim:
into ``guardrail_response`` and from there onto OTEL guardrail spans."""
@staticmethod
def _logged_response(request_data: dict) -> object:
entries = request_data["litellm_metadata"]["standard_logging_guardrail_information"]
def _logged_response(request_data: dict[str, object]) -> object:
metadata = request_data["litellm_metadata"]
assert isinstance(metadata, dict)
entries = metadata["standard_logging_guardrail_information"]
assert len(entries) == 1
return entries[0]["guardrail_response"]
@staticmethod
def _request() -> dict[str, object]:
return {
"model": "gpt-4.1-mini",
"input": "SECRET_PROMPT",
"messages": [{"role": "user", "content": "SECRET_PROMPT"}],
"litellm_metadata": {},
}
@pytest.mark.asyncio
async def test_pre_call_hook_returning_request_logs_allow(self):
class PassthroughGuardrail(CustomGuardrail):
@log_guardrail_information
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: UserAPIKeyAuth,
cache: object,
data: dict[str, object],
call_type: str,
) -> dict[str, object]:
return data
guardrail = PassthroughGuardrail(guardrail_name="g")
data = {"model": "gpt-4.1-mini", "input": "SECRET_PROMPT", "litellm_metadata": {}}
await guardrail.async_pre_call_hook(
data = self._request()
await PassthroughGuardrail(guardrail_name="g").async_pre_call_hook(
user_api_key_dict=UserAPIKeyAuth(), cache=None, data=data, call_type="aresponses"
)
@ -2538,14 +2553,41 @@ class TestPreCallHookResponseIsNotLoggedVerbatim:
async def test_pre_call_hook_returning_modified_copy_logs_mask(self):
class MaskingGuardrail(CustomGuardrail):
@log_guardrail_information
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: UserAPIKeyAuth,
cache: object,
data: dict[str, object],
call_type: str,
) -> dict[str, object]:
return {**data, "input": "[MASKED]"}
guardrail = MaskingGuardrail(guardrail_name="g")
data = {"model": "gpt-4.1-mini", "input": "SECRET_PROMPT", "litellm_metadata": {}}
await guardrail.async_pre_call_hook(
data = self._request()
await MaskingGuardrail(guardrail_name="g").async_pre_call_hook(
user_api_key_dict=UserAPIKeyAuth(), cache=None, data=data, call_type="aresponses"
)
assert self._logged_response(data) == "mask"
@pytest.mark.asyncio
async def test_pre_call_hook_mutating_request_in_place_logs_mask(self):
class InPlaceMaskingGuardrail(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]:
messages = data["messages"]
assert isinstance(messages, list)
messages[0]["content"] = "[MASKED]"
return data
data = self._request()
await InPlaceMaskingGuardrail(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"