diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index 0f4ebbd4880..d71884bf17c 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -504,13 +504,9 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): "This may indicate a missing caller update." ) request_data = {} - # Store pii_tokens in metadata to avoid leaking to LLM providers. - # Providers like Anthropic reject unknown top-level fields. - 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._get_or_create_request_scoped_pii_tokens( + request_data + ) # Append a sequential number to make each token unique # per request, so unmasking maps back to the correct @@ -935,6 +931,56 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): ) return response + def _get_or_create_request_scoped_pii_tokens( + self, request_data: Dict[str, Any] + ) -> Dict[str, str]: + """ + Return request-scoped token map and ensure it is anchored in all + metadata locations used in pre/post guardrail hooks. + """ + metadata = request_data.get("metadata") + if not isinstance(metadata, dict): + metadata = {} + request_data["metadata"] = metadata + + pii_tokens = metadata.get("pii_tokens") + if not isinstance(pii_tokens, dict): + pii_tokens = {} + + metadata["pii_tokens"] = pii_tokens + + litellm_metadata = request_data.get("litellm_metadata") + if isinstance(litellm_metadata, dict): + litellm_metadata["pii_tokens"] = pii_tokens + + return cast(Dict[str, str], pii_tokens) + + def _get_request_scoped_pii_tokens( + self, request_data: Optional[Dict[str, Any]] + ) -> Dict[str, str]: + """ + Read request-scoped pii token map with fallback across known metadata paths. + """ + if not request_data: + return {} + + metadata = request_data.get("metadata") + if isinstance(metadata, dict): + pii_tokens = metadata.get("pii_tokens") + if isinstance(pii_tokens, dict): + litellm_metadata = request_data.get("litellm_metadata") + if isinstance(litellm_metadata, dict): + litellm_metadata["pii_tokens"] = pii_tokens + return cast(Dict[str, str], pii_tokens) + + litellm_metadata = request_data.get("litellm_metadata") + if isinstance(litellm_metadata, dict): + pii_tokens = litellm_metadata.get("pii_tokens") + if isinstance(pii_tokens, dict): + return cast(Dict[str, str], pii_tokens) + + return {} + @staticmethod def _unmask_pii_text(text: str, pii_tokens: Dict[str, str]) -> str: """ @@ -980,11 +1026,10 @@ 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_request_scoped_pii_tokens(request_data) if not pii_tokens and mode == "unmask": verbose_proxy_logger.debug( - "No pii_tokens in metadata for Anthropic response unmask" + "No pii_tokens found for Anthropic response unmask" ) presidio_config = self.get_presidio_settings_from_request_data( request_data or {} @@ -1022,11 +1067,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_request_scoped_pii_tokens(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 — nothing to unmask" ) presidio_config = self.get_presidio_settings_from_request_data( request_data or {} @@ -1253,11 +1297,10 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): yield chunk return - metadata = (request_data.get("metadata") or {}) if request_data else {} - pii_tokens = metadata.get("pii_tokens", {}) + pii_tokens = self._get_request_scoped_pii_tokens(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 found in request_data for streaming unmask path" ) if not (self.output_parse_pii and pii_tokens): async for chunk in response: @@ -1317,13 +1360,28 @@ 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_request_scoped_pii_tokens(request_data) new_texts = [] - if input_type == "response" and pii_tokens: - for text in texts: - new_texts.append(self._unmask_pii_text(text, pii_tokens)) + if input_type == "response": + if pii_tokens: + for text in texts: + new_texts.append(self._unmask_pii_text(text, pii_tokens)) + elif self.output_parse_pii: + verbose_proxy_logger.debug( + "output_parse_pii enabled but no pii_tokens found on response path; " + "leaving response text unchanged" + ) + new_texts = list(texts) + else: + for text in texts: + modified_text = await self.check_pii( + text=text, + output_parse_pii=self.output_parse_pii, + presidio_config=None, + request_data=request_data or {}, + ) + new_texts.append(modified_text) else: for text in texts: modified_text = await self.check_pii( 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 32a8c1b1070..f860bab3813 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py @@ -2161,6 +2161,62 @@ async def test_apply_guardrail_unmask_on_response(): assert result["texts"][0] == "Hello John Smith, your number is 555-123-4567." +@pytest.mark.asyncio +async def test_apply_guardrail_unmask_reads_tokens_from_litellm_metadata(): + """ + Regression test: if pii_tokens are present on litellm_metadata (but not metadata), + response unmasking should still work. + """ + guardrail = _OPTIONAL_PresidioPIIMasking( + guardrail_name="test_presidio", + output_parse_pii=True, + mock_testing=True, + ) + + request_data = { + "model": "gpt-4o", + "litellm_metadata": { + "pii_tokens": { + "": "Jane Doe", + } + }, + } + + result = await guardrail.apply_guardrail( + inputs={"texts": ["Hello "]}, + request_data=request_data, + input_type="response", + ) + + assert result["texts"][0] == "Hello Jane Doe" + + +@pytest.mark.asyncio +async def test_apply_guardrail_response_no_tokens_does_not_remask(): + """ + Regression test: for output_parse_pii response path with no tokens, + text should be returned unchanged (no re-masking pass). + """ + guardrail = _OPTIONAL_PresidioPIIMasking( + guardrail_name="test_presidio", + output_parse_pii=True, + mock_testing=True, + ) + + async def fail_if_called(*args, **kwargs): + raise AssertionError("check_pii should not run on response path without tokens") + + guardrail.check_pii = fail_if_called + + result = await guardrail.apply_guardrail( + inputs={"texts": ["Nice to meet you, !"]}, + request_data={"model": "gpt-4o", "metadata": {}}, + input_type="response", + ) + + assert result["texts"][0] == "Nice to meet you, !" + + @pytest.mark.asyncio async def test_apply_guardrail_masks_on_request(): """