diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index ce32ebf54f8..3d2ab334402 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -100,6 +100,16 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): self.mock_redacted_text = mock_redacted_text self.output_parse_pii = output_parse_pii or False self.apply_to_output = apply_to_output + + # When output_parse_pii or apply_to_output is enabled, the guardrail must + # also run on post_call to unmask/mask the response. Expand the event_hook + # so should_run_guardrail returns True for both pre_call and post_call. + if (self.output_parse_pii or self.apply_to_output) and not logging_only: + current_hook = self.event_hook + if isinstance(current_hook, str) and current_hook == "pre_call": + self.event_hook = ["pre_call", "post_call"] + elif isinstance(current_hook, list) and "post_call" not in current_hook: + self.event_hook = current_hook + ["post_call"] self.pii_entities_config: Dict[Union[PiiEntityType, str], PiiAction] = ( pii_entities_config or {} ) @@ -489,9 +499,13 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): "This may indicate a missing caller update." ) request_data = {} - if "pii_tokens" not in request_data: - request_data["pii_tokens"] = {} - pii_tokens = request_data["pii_tokens"] + # Store pii_tokens in metadata to avoid leaking to LLM providers. + # Providers like Anthropic reject unknown top-level fields. + if "metadata" not in request_data: + request_data["metadata"] = {} + if "pii_tokens" not in request_data["metadata"]: + request_data["metadata"]["pii_tokens"] = {} + pii_tokens = request_data["metadata"]["pii_tokens"] # Always append a UUID to ensure the replacement token is unique to this request and session. # This prevents collisions where the LLM might hallucinate a generic token like [PHONE_NUMBER]. @@ -544,8 +558,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): filtered_results: List[PresidioAnalyzeResponseItem] = [] deny_list_strings = [ - getattr(x, "value", str(x)) - for x in self.presidio_entities_deny_list + getattr(x, "value", str(x)) for x in self.presidio_entities_deny_list ] for item in analyze_results: entity_type = item.get("entity_type") @@ -937,10 +950,11 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): Helper to recursively process a ModelResponse for PII. Handles all choices and tool calls. """ - pii_tokens = request_data.get("pii_tokens", {}) if request_data else {} + metadata = request_data.get("metadata", {}) if request_data else {} + pii_tokens = metadata.get("pii_tokens", {}) if not pii_tokens and mode == "unmask": verbose_proxy_logger.debug( - "No pii_tokens found in request_data — 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 {} @@ -1099,10 +1113,11 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): return # --- PII unmasking path (output_parse_pii=True) --- - pii_tokens = request_data.get("pii_tokens", {}) if request_data else {} + metadata = request_data.get("metadata", {}) if request_data else {} + pii_tokens = metadata.get("pii_tokens", {}) if not pii_tokens and request_data: verbose_proxy_logger.debug( - "No pii_tokens in request_data for streaming unmask path" + "No pii_tokens in request_data['metadata'] for streaming unmask path" ) if not (self.output_parse_pii and pii_tokens): async for chunk in response: 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 76f9c39acd0..e336f2f1bf8 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py @@ -1608,3 +1608,106 @@ async def test_anonymize_text_http_error_status(): output_parse_pii=False, masked_entity_count={}, ) + + +@pytest.mark.asyncio +async def test_pii_tokens_stored_in_metadata_not_top_level(presidio_guardrail): + """ + Regression test: pii_tokens must be stored in data['metadata']['pii_tokens'], + NOT in data['pii_tokens']. Storing at the top level leaks the field to LLM + providers like Anthropic, which reject unknown fields with + 'pii_tokens: Extra inputs are not permitted'. + """ + guardrail = _OPTIONAL_PresidioPIIMasking( + mock_testing=True, + output_parse_pii=True, + pii_entities_config={ + PiiEntityType.PERSON: PiiAction.MASK, + PiiEntityType.PHONE_NUMBER: PiiAction.MASK, + }, + ) + + mock_user_api_key = UserAPIKeyAuth(api_key="test-key") + mock_cache = DualCache() + + test_data = { + "messages": [ + {"role": "user", "content": "My name is John and my phone is 555-123-4567"} + ], + "model": "claude-haiku-4-5-20251001", + "metadata": {}, + } + + async def mock_check_pii(text, output_parse_pii, presidio_config, request_data): + # Simulate PII masking with token storage (mimics real anonymize_text behavior) + import uuid + + if request_data is not None and output_parse_pii: + if "metadata" not in request_data: + request_data["metadata"] = {} + if "pii_tokens" not in request_data["metadata"]: + request_data["metadata"]["pii_tokens"] = {} + pii_tokens = request_data["metadata"]["pii_tokens"] + token = f"_{str(uuid.uuid4())[:12]}" + pii_tokens[token] = "John" + text = text.replace("John", token) + return text + + guardrail.check_pii = mock_check_pii + + result = await guardrail.async_pre_call_hook( + user_api_key_dict=mock_user_api_key, + cache=mock_cache, + data=test_data, + call_type="completion", + ) + + # pii_tokens must NOT be at the top level of data (would leak to providers) + assert "pii_tokens" not in result, ( + "pii_tokens must not be a top-level key in request data — " + "it would leak to LLM providers and cause 'Extra inputs are not permitted' errors" + ) + + # pii_tokens must be inside metadata (safe from provider leakage) + assert "metadata" in result + assert "pii_tokens" in result["metadata"] + assert len(result["metadata"]["pii_tokens"]) > 0 + + +@pytest.mark.asyncio +async def test_pii_tokens_in_metadata_used_for_unmasking(): + """ + Regression test: _process_response_for_pii must read pii_tokens from + data['metadata']['pii_tokens'] and correctly unmask the response. + """ + guardrail = _OPTIONAL_PresidioPIIMasking( + mock_testing=True, + output_parse_pii=True, + ) + + token_key = "_abc123def456" + request_data = { + "model": "claude-haiku-4-5-20251001", + "metadata": {"pii_tokens": {token_key: "John"}}, + } + + response = ModelResponse( + choices=[ + Choices( + message=Message( + role="assistant", + content=f"Hello {token_key}, how can I help you?", + ), + index=0, + finish_reason="stop", + ) + ] + ) + + await guardrail._process_response_for_pii( + response=response, + request_data=request_data, + mode="unmask", + ) + + assert response.choices[0].message.content == "Hello John, how can I help you?"