From 49b49f5c5291cbad66bd343e37e13eec0100552f Mon Sep 17 00:00:00 2001 From: shivam Date: Fri, 4 Sep 2026 03:01:34 +0000 Subject: [PATCH] fix(guardrails): snapshot pre_call request before the hook so in-place edits log as mask Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/integrations/custom_guardrail.py | 63 +++++++++--------- .../guardrail_hooks/azure/prompt_shield.py | 2 - .../integrations/test_custom_guardrail.py | 66 +++++++++++++++---- 3 files changed, 85 insertions(+), 46 deletions(-) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 4eca77b77c5..e8bc0698223 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -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) diff --git a/litellm/proxy/guardrails/guardrail_hooks/azure/prompt_shield.py b/litellm/proxy/guardrails/guardrail_hooks/azure/prompt_shield.py index b54be6526ea..de9618a44a1 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/azure/prompt_shield.py +++ b/litellm/proxy/guardrails/guardrail_hooks/azure/prompt_shield.py @@ -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( diff --git a/tests/test_litellm/integrations/test_custom_guardrail.py b/tests/test_litellm/integrations/test_custom_guardrail.py index 6c3e42d95e7..8a53cbbaef3 100644 --- a/tests/test_litellm/integrations/test_custom_guardrail.py +++ b/tests/test_litellm/integrations/test_custom_guardrail.py @@ -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"