diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 8a976a966a6..39adea30828 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -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} diff --git a/litellm/proxy/guardrails/guardrail_hooks/azure/prompt_shield.py b/litellm/proxy/guardrails/guardrail_hooks/azure/prompt_shield.py index de9618a44a1..a0724b75ec7 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/azure/prompt_shield.py +++ b/litellm/proxy/guardrails/guardrail_hooks/azure/prompt_shield.py @@ -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, diff --git a/tests/test_litellm/integrations/test_custom_guardrail.py b/tests/test_litellm/integrations/test_custom_guardrail.py index bb4822eae57..2fd5fa76d8e 100644 --- a/tests/test_litellm/integrations/test_custom_guardrail.py +++ b/tests/test_litellm/integrations/test_custom_guardrail.py @@ -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"] = [""] + 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"