Merge remote-tracking branch 'origin/litellm_internal_staging' into litellm_/release-version-bump-5f3482

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}

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,

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"] = ["<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"