From 9d7fc90c545bb017509b0efdcb5c3aced052f9e4 Mon Sep 17 00:00:00 2001 From: shivam Date: Tue, 19 May 2026 19:09:29 -0700 Subject: [PATCH] fix(presidio): restore PII unmasking in streaming responses When output_parse_pii is enabled with presidio_filter_scope=both, a separate apply_to_output callback was re-masking the stream after unmask ran. Skip registering that callback when output_parse_pii is on, mirror pii_tokens across metadata containers, and passthrough apply_to_output streaming when tokens exist. Co-authored-by: Cursor --- .../guardrails/guardrail_hooks/presidio.py | 75 ++++++++-- .../guardrails/guardrail_initializers.py | 2 +- .../guardrail_hooks/test_presidio.py | 130 +++++++++++++++++- 3 files changed, 190 insertions(+), 17 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index fc414ab7b54..954bae184a4 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -482,6 +482,40 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): ) return redacted_text["text"] + def _ensure_pii_tokens(self, request_data: dict) -> Dict[str, str]: + """ + Return the shared pii_tokens map for this request, stored on both + metadata and litellm_metadata so post-call unmasking works regardless + of which container downstream hooks read. + """ + tokens: Optional[Dict[str, str]] = None + for meta_key in ("metadata", "litellm_metadata"): + meta = request_data.get(meta_key) + if isinstance(meta, dict) and isinstance(meta.get("pii_tokens"), dict): + tokens = meta["pii_tokens"] + break + if tokens is None: + tokens = {} + for meta_key in ("metadata", "litellm_metadata"): + if not isinstance(request_data.get(meta_key), dict): + request_data[meta_key] = {} + request_data[meta_key]["pii_tokens"] = tokens + return tokens + + @staticmethod + def _get_pii_tokens_from_request_data( + request_data: Optional[Dict], + ) -> Dict[str, str]: + if not request_data: + return {} + for meta_key in ("metadata", "litellm_metadata"): + meta = request_data.get(meta_key) + if isinstance(meta, dict): + tokens = meta.get("pii_tokens") + if isinstance(tokens, dict) and tokens: + return tokens + return {} + def _finalize_presidio_anonymize_numbered_tokens( self, text: str, @@ -501,11 +535,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): "This may indicate a missing caller update." ) request_data = {} - if not request_data.get("metadata"): - request_data["metadata"] = {} - if "pii_tokens" not in request_data["metadata"]: - request_data["metadata"]["pii_tokens"] = {} - pii_tokens = request_data["metadata"]["pii_tokens"] + pii_tokens = self._ensure_pii_tokens(request_data) # Assign sequence numbers in forward (left-to-right) order so # that is the first entity in the text, etc. @@ -1001,8 +1031,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): Process an Anthropic native message dict for PII masking/unmasking. Handles content blocks with type == "text". """ - metadata = (request_data.get("metadata") or {}) if request_data else {} - pii_tokens = metadata.get("pii_tokens", {}) + pii_tokens = self._get_pii_tokens_from_request_data(request_data) if not pii_tokens and mode == "unmask": verbose_proxy_logger.debug( "No pii_tokens in metadata for Anthropic response unmask" @@ -1043,11 +1072,10 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): Helper to recursively process a ModelResponse for PII. Handles all choices and tool calls. """ - metadata = (request_data.get("metadata") or {}) if request_data else {} - pii_tokens = metadata.get("pii_tokens", {}) + pii_tokens = self._get_pii_tokens_from_request_data(request_data) if not pii_tokens and mode == "unmask": verbose_proxy_logger.debug( - "No pii_tokens found in request_data['metadata'] — nothing to unmask" + "No pii_tokens found in request_data metadata — nothing to unmask" ) presidio_config = self.get_presidio_settings_from_request_data( request_data or {} @@ -1292,17 +1320,21 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): The base class declares ModelResponseStream only. """ if self.apply_to_output: + pii_tokens = self._get_pii_tokens_from_request_data(request_data) + if pii_tokens: + async for chunk in response: + yield chunk + return async for chunk in self._stream_apply_output_masking( response, request_data ): yield chunk return - metadata = (request_data.get("metadata") or {}) if request_data else {} - pii_tokens = metadata.get("pii_tokens", {}) + pii_tokens = self._get_pii_tokens_from_request_data(request_data) if not pii_tokens and request_data: verbose_proxy_logger.debug( - "No pii_tokens in request_data['metadata'] for streaming unmask path" + "No pii_tokens in request_data for streaming unmask path" ) if not (self.output_parse_pii and pii_tokens): async for chunk in response: @@ -1362,8 +1394,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): # When input_type is "response" and pii_tokens are available, # unmask the text instead of masking it. - metadata = (request_data.get("metadata") or {}) if request_data else {} - pii_tokens = metadata.get("pii_tokens", {}) + pii_tokens = self._get_pii_tokens_from_request_data(request_data) new_texts = [] if input_type == "response" and pii_tokens: @@ -1394,3 +1425,17 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): self.presidio_entities_deny_list = ( litellm_params.presidio_entities_deny_list ) + if litellm_params.output_parse_pii is not None: + self.output_parse_pii = litellm_params.output_parse_pii or False + if (self.output_parse_pii or self.apply_to_output) and not getattr( + self, "logging_only", False + ): + current_hook = self.event_hook + if isinstance(current_hook, str) and current_hook != "post_call": + self.event_hook = cast( + List[GuardrailEventHooks], [current_hook, "post_call"] + ) + elif isinstance(current_hook, list) and "post_call" not in current_hook: + self.event_hook = cast( + List[GuardrailEventHooks], current_hook + ["post_call"] + ) diff --git a/litellm/proxy/guardrails/guardrail_initializers.py b/litellm/proxy/guardrails/guardrail_initializers.py index 109f2237165..ab24fdb5461 100644 --- a/litellm/proxy/guardrails/guardrail_initializers.py +++ b/litellm/proxy/guardrails/guardrail_initializers.py @@ -112,7 +112,7 @@ def initialize_presidio(litellm_params: LitellmParams, guardrail: Guardrail): event_hook=GuardrailEventHooks.post_call.value, ) - if run_output: + if run_output and not litellm_params.output_parse_pii: output_callback = _make_presidio_callback( apply_to_output=True, event_hook=GuardrailEventHooks.post_call.value, diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py index 565bf83c6a2..f235d869a82 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py @@ -21,7 +21,14 @@ from litellm.proxy.guardrails.guardrail_hooks.presidio import ( ) from litellm.exceptions import GuardrailRaisedException from litellm.types.guardrails import LitellmParams, PiiAction, PiiEntityType -from litellm.types.utils import Choices, Message, ModelResponse +from litellm.types.utils import ( + Choices, + Delta, + Message, + ModelResponse, + ModelResponseStream, + StreamingChoices, +) def _make_mock_session_iterator( @@ -797,6 +804,61 @@ async def test_presidio_filter_scope_initializer(monkeypatch): assert any(not c.apply_to_output for c in created) assert any(c.apply_to_output for c in created) + # both + output_parse_pii -> unmask only; no apply_to_output callback + created.clear() + params_both_unmask = LitellmParams( + guardrail="presidio", + mode="pre_call", + presidio_filter_scope="both", + output_parse_pii=True, + ) + cb = initialize_presidio(params_both_unmask, guardrail_dict) + assert len(created) == 2 + assert all(not c.apply_to_output for c in created) + + +@pytest.mark.asyncio +async def test_apply_to_output_streaming_skips_remask_when_pii_tokens_present(): + """ + When input was masked with numbered tokens, the apply_to_output streaming + hook must not re-mask the already-unmasked response. + """ + guardrail = _OPTIONAL_PresidioPIIMasking( + guardrail_name="test_presidio", + apply_to_output=True, + event_hook="post_call", + mock_testing=True, + ) + + request_data = { + "metadata": { + "pii_tokens": {"": "shivam@uni.minerva.edu"}, + } + } + + async def _fake_stream(): + yield ModelResponseStream( + choices=[ + StreamingChoices( + delta=Delta(content="Your email is shivam@uni.minerva.edu"), + finish_reason="stop", + index=0, + ) + ] + ) + + chunks = [] + async for chunk in guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), + response=_fake_stream(), + request_data=request_data, + ): + chunks.append(chunk) + + assert len(chunks) == 1 + assert "shivam@uni.minerva.edu" in chunks[0].choices[0].delta.content + assert "": "John Smith"}, + }, + "litellm_metadata": {"user_api_key_user_id": "user-1"}, + } + response = ModelResponse( + choices=[ + Choices( + message=Message( + role="assistant", + content="Hello , how can I help?", + ), + index=0, + finish_reason="stop", + ) + ] + ) + + proxy_logging = ProxyLogging(user_api_key_cache=None) + result = await proxy_logging.post_call_success_hook( + data=data, + response=response, + user_api_key_dict=UserAPIKeyAuth( + api_key="test-key", request_route="/chat/completions" + ), + ) + + assert result.choices[0].message.content == "Hello John Smith, how can I help?" + + +def test_pii_tokens_mirrored_to_litellm_metadata(): + """pii_tokens must be readable from either metadata container after masking.""" + guardrail = _OPTIONAL_PresidioPIIMasking(mock_testing=True) + request_data: dict = {"litellm_metadata": {"user_api_key_user_id": "u1"}} + tokens = guardrail._ensure_pii_tokens(request_data) + tokens[""] = "John Smith" + + assert request_data["metadata"]["pii_tokens"][""] == "John Smith" + assert request_data["litellm_metadata"]["pii_tokens"][""] == "John Smith" + assert guardrail._get_pii_tokens_from_request_data(request_data)[""] == ( + "John Smith" + ) + + @pytest.mark.asyncio async def test_apply_guardrail_masks_on_request(): """