diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index ce32ebf54f8..ddeba2100c3 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -36,7 +36,6 @@ from litellm.types.utils import GenericGuardrailAPIInputs if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj -from litellm._uuid import uuid from litellm.caching.caching import DualCache from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException from litellm.integrations.custom_guardrail import ( @@ -100,6 +99,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 != "post_call": + self.event_hook = [current_hook, "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 {} ) @@ -475,13 +484,15 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): new_text = text if redacted_text is not None: verbose_proxy_logger.debug("redacted_text: %s", redacted_text) - for item in redacted_text["items"]: + # Process items in reverse order by start position so that + # replacing later spans first does not shift earlier coordinates. + for item in sorted( + redacted_text["items"], key=lambda x: x["start"], reverse=True + ): start = item["start"] end = item["end"] replacement = item["text"] # replacement token if item["operator"] == "replace" and output_parse_pii is True: - # check if token in dict - # if exists, add a uuid to the replacement token for swapping back to the original text in llm response output parsing if request_data is None: verbose_proxy_logger.warning( "Presidio anonymize_text called without request_data — " @@ -489,17 +500,28 @@ 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 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"] - # 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]. - replacement = f"{replacement}_{str(uuid.uuid4())[:12]}" + # Append a sequential number to make each token unique + # per request, so unmasking maps back to the correct + # original value. Format: , + # This is LLM-friendly and degrades gracefully if the + # LLM doesn't echo the token verbatim. + seq = len(pii_tokens) + 1 + if replacement.endswith(">"): + replacement = f"{replacement[:-1]}_{seq}>" + else: + replacement = f"{replacement}_{seq}" - pii_tokens[replacement] = new_text[ - start:end - ] # get text it'll replace + # Use ORIGINAL text (not new_text) since start/end + # reference the original text's coordinates. + pii_tokens[replacement] = text[start:end] new_text = new_text[:start] + replacement + new_text[end:] entity_type = item.get("entity_type", None) @@ -507,12 +529,13 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): masked_entity_count[entity_type] = ( masked_entity_count.get(entity_type, 0) + 1 ) - # When output_parse_pii is True, new_text contains UUID-suffixed - # tokens that match the keys in pii_tokens. Returning - # redacted_text["text"] (Presidio's original output) would send - # un-suffixed tokens to the LLM, making unmasking impossible. + # When output_parse_pii is True, new_text contains sequentially + # numbered tokens (e.g. ) that match the keys + # in pii_tokens. Returning redacted_text["text"] (Presidio's + # original output) would send un-numbered tokens to the LLM, + # making unmasking impossible. # When output_parse_pii is False, new_text == redacted_text["text"] - # because no UUID suffix is appended. + # because no suffix is appended. return new_text else: raise Exception("Invalid anonymizer response: received None") @@ -544,8 +567,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") @@ -884,6 +906,10 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): ) if self.apply_to_output is True: + if self._is_anthropic_message_response(response): + return await self._process_anthropic_response_for_pii( + response=response, request_data=data, mode="mask" + ) return await self._mask_output_response( response=response, request_data=data ) @@ -899,6 +925,10 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): request_data=data, mode="unmask", ) + elif self._is_anthropic_message_response(response): + await self._process_anthropic_response_for_pii( + response=response, request_data=data, mode="unmask" + ) return response @staticmethod @@ -927,6 +957,57 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): break return text + @staticmethod + def _is_anthropic_message_response(response: Any) -> bool: + """Check if the response is an Anthropic native message dict.""" + return ( + isinstance(response, dict) + and response.get("type") == "message" + and isinstance(response.get("content"), list) + ) + + async def _process_anthropic_response_for_pii( + self, + response: dict, + request_data: dict, + mode: Literal["mask", "unmask"], + ) -> dict: + """ + 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", {}) + if not pii_tokens and mode == "unmask": + verbose_proxy_logger.debug( + "No pii_tokens in metadata for Anthropic response unmask" + ) + presidio_config = self.get_presidio_settings_from_request_data( + request_data or {} + ) + + content = response.get("content") + if not isinstance(content, list): + return response + + for block in content: + if not isinstance(block, dict) or block.get("type") != "text": + continue + text_value = block.get("text") + if text_value is None: + continue + if mode == "unmask": + block["text"] = self._unmask_pii_text(text_value, pii_tokens) + elif mode == "mask": + block["text"] = await self.check_pii( + text=text_value, + output_parse_pii=False, + presidio_config=presidio_config, + request_data=request_data, + ) + + return response + async def _process_response_for_pii( self, response: ModelResponse, @@ -937,10 +1018,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") or {}) 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 {} @@ -1045,7 +1127,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): user_api_key_dict: UserAPIKeyAuth, response: Any, request_data: dict, - ) -> AsyncGenerator[ModelResponseStream, None]: + ) -> AsyncGenerator[Union[ModelResponseStream, bytes], None]: """ Process streaming response chunks to unmask PII tokens when needed. """ @@ -1062,8 +1144,20 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): async for chunk in response: if isinstance(chunk, ModelResponseStream): all_chunks.append(chunk) + elif isinstance(chunk, bytes): + # Anthropic native SSE: pass through as-is + yield chunk # type: ignore[misc] + continue if not all_chunks: + # All chunks were Anthropic native SSE bytes — output + # masking cannot be applied to raw bytes. Log a warning + # so operators know PII masking was skipped for this stream. + verbose_proxy_logger.warning( + "Presidio apply_to_output: streaming response contained only " + "bytes chunks (Anthropic native SSE). Output PII masking was " + "skipped for this response." + ) return assembled_model_response = stream_chunk_builder( @@ -1099,10 +1193,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") or {}) 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: @@ -1114,6 +1209,10 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): async for chunk in response: if isinstance(chunk, ModelResponseStream): remaining_chunks.append(chunk) + elif isinstance(chunk, bytes): + # Anthropic native SSE: pass through as-is + yield chunk # type: ignore[misc] + continue if not remaining_chunks: return @@ -1191,15 +1290,24 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): """ texts = inputs.get("texts", []) + # 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", {}) + new_texts = [] - 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) + if input_type == "response" and pii_tokens: + for text in texts: + new_texts.append(self._unmask_pii_text(text, pii_tokens)) + 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) inputs["texts"] = new_texts return inputs diff --git a/tests/guardrails_tests/test_presidio_pii.py b/tests/guardrails_tests/test_presidio_pii.py index 0d730288e63..eda0c7bb5b5 100644 --- a/tests/guardrails_tests/test_presidio_pii.py +++ b/tests/guardrails_tests/test_presidio_pii.py @@ -1,20 +1,19 @@ import sys import os -import io, asyncio import pytest -import time from litellm import mock_completion -from unittest.mock import MagicMock, AsyncMock, patch +from unittest.mock import patch + sys.path.insert(0, os.path.abspath("../..")) import litellm -from litellm.proxy.guardrails.guardrail_hooks.presidio import _OPTIONAL_PresidioPIIMasking, PresidioPerRequestConfig +from litellm.proxy.guardrails.guardrail_hooks.presidio import ( + _OPTIONAL_PresidioPIIMasking, + PresidioPerRequestConfig, +) from litellm.types.guardrails import PiiEntityType, PiiAction from litellm.proxy._types import UserAPIKeyAuth from litellm.caching.caching import DualCache from litellm.exceptions import BlockedPiiEntityError -from litellm.types.utils import CallTypes as LitellmCallTypes - - @pytest.mark.asyncio @@ -26,42 +25,37 @@ async def test_presidio_with_entities_config(): PiiEntityType.CREDIT_CARD: PiiAction.MASK, PiiEntityType.EMAIL_ADDRESS: PiiAction.MASK, } - + presidio_guardrail = _OPTIONAL_PresidioPIIMasking( pii_entities_config=pii_entities_config, presidio_analyzer_api_base=os.environ.get("PRESIDIO_ANALYZER_API_BASE"), - presidio_anonymizer_api_base=os.environ.get("PRESIDIO_ANONYMIZER_API_BASE") + presidio_anonymizer_api_base=os.environ.get("PRESIDIO_ANONYMIZER_API_BASE"), ) - + # Test text with different PII types test_text = "My credit card number is 4111-1111-1111-1111, my email is test@example.com, and my phone is 555-123-4567" - + # Test the analyze request configuration analyze_request = presidio_guardrail._get_presidio_analyze_request_payload( - text=test_text, - presidio_config=None, - request_data={} + text=test_text, presidio_config=None, request_data={} ) - + # Verify entities were passed correctly assert "entities" in analyze_request assert set(analyze_request["entities"]) == set(pii_entities_config.keys()) - + # Test the check_pii method - this will call the actual Presidio API redacted_text = await presidio_guardrail.check_pii( - text=test_text, - output_parse_pii=True, - presidio_config=None, - request_data={} + text=test_text, output_parse_pii=True, presidio_config=None, request_data={} ) - + # Verify PII has been masked/replaced/redacted in the result assert "4111-1111-1111-1111" not in redacted_text assert "test@example.com" not in redacted_text # Since this entity is not in the config, it should not be masked assert "555-123-4567" in redacted_text - + # The specific replacements will vary based on Presidio's implementation print(f"Redacted text: {redacted_text}") @@ -73,10 +67,12 @@ async def test_presidio_apply_guardrail(): presidio_guardrail = _OPTIONAL_PresidioPIIMasking( pii_entities_config={}, presidio_analyzer_api_base=os.environ.get("PRESIDIO_ANALYZER_API_BASE"), - presidio_anonymizer_api_base=os.environ.get("PRESIDIO_ANONYMIZER_API_BASE") + presidio_anonymizer_api_base=os.environ.get("PRESIDIO_ANONYMIZER_API_BASE"), ) - test_text = "My credit card number is 4111-1111-1111-1111 and my email is test@example.com" + test_text = ( + "My credit card number is 4111-1111-1111-1111 and my email is test@example.com" + ) response = await presidio_guardrail.apply_guardrail( inputs={"texts": [test_text]}, request_data={}, @@ -91,6 +87,7 @@ async def test_presidio_apply_guardrail(): assert "4111-1111-1111-1111" not in modified_text assert "test@example.com" not in modified_text + @pytest.mark.asyncio async def test_presidio_with_blocked_entities(): """Test for Presidio guardrail with blocked entities - requires actual Presidio API""" @@ -100,36 +97,33 @@ async def test_presidio_with_blocked_entities(): PiiEntityType.CREDIT_CARD: PiiAction.BLOCK, # This entity should cause a block PiiEntityType.EMAIL_ADDRESS: PiiAction.MASK, # This entity should be masked } - + presidio_guardrail = _OPTIONAL_PresidioPIIMasking( pii_entities_config=pii_entities_config, presidio_analyzer_api_base=os.environ.get("PRESIDIO_ANALYZER_API_BASE"), - presidio_anonymizer_api_base=os.environ.get("PRESIDIO_ANONYMIZER_API_BASE") + presidio_anonymizer_api_base=os.environ.get("PRESIDIO_ANONYMIZER_API_BASE"), ) - + # Test text with blocked PII type - test_text = "My credit card number is 4111-1111-1111-1111 and my email is test@example.com" - + test_text = ( + "My credit card number is 4111-1111-1111-1111 and my email is test@example.com" + ) + # Verify the analyze request configuration analyze_request = presidio_guardrail._get_presidio_analyze_request_payload( - text=test_text, - presidio_config=None, - request_data={} + text=test_text, presidio_config=None, request_data={} ) - + # Verify entities were passed correctly assert "entities" in analyze_request assert set(analyze_request["entities"]) == set(pii_entities_config.keys()) - + # Test that BlockedPiiEntityError is raised when check_pii is called with pytest.raises(BlockedPiiEntityError) as excinfo: await presidio_guardrail.check_pii( - text=test_text, - output_parse_pii=True, - presidio_config=None, - request_data={} + text=test_text, output_parse_pii=True, presidio_config=None, request_data={} ) - + # Verify the error contains the correct entity type assert excinfo.value.entity_type == PiiEntityType.CREDIT_CARD assert excinfo.value.guardrail_name == presidio_guardrail.guardrail_name @@ -143,37 +137,40 @@ async def test_presidio_pre_call_hook_with_blocked_entities(): PiiEntityType.CREDIT_CARD: PiiAction.BLOCK, # This entity should cause a block PiiEntityType.EMAIL_ADDRESS: PiiAction.MASK, # This entity should be masked } - + presidio_guardrail = _OPTIONAL_PresidioPIIMasking( pii_entities_config=pii_entities_config, presidio_analyzer_api_base=os.environ.get("PRESIDIO_ANALYZER_API_BASE"), - presidio_anonymizer_api_base=os.environ.get("PRESIDIO_ANONYMIZER_API_BASE") + presidio_anonymizer_api_base=os.environ.get("PRESIDIO_ANONYMIZER_API_BASE"), ) - + # Create a sample chat completion request with PII data data = { "messages": [ {"role": "system", "content": "You are a helpful assistant."}, - {"role": "user", "content": "My credit card is 4111-1111-1111-1111 and my email is test@example.com."} + { + "role": "user", + "content": "My credit card is 4111-1111-1111-1111 and my email is test@example.com.", + }, ], - "model": "gpt-3.5-turbo" + "model": "gpt-3.5-turbo", } - + # Mock objects needed for the pre-call hook user_api_key_dict = UserAPIKeyAuth(api_key="test_key") cache = DualCache() - + # Call the pre-call hook and expect BlockedPiiEntityError with pytest.raises(BlockedPiiEntityError) as excinfo: await presidio_guardrail.async_pre_call_hook( user_api_key_dict=user_api_key_dict, cache=cache, data=data, - call_type="completion" + call_type="completion", ) - + print(f"got error: {excinfo}") - + # Verify the error contains the correct entity type assert excinfo.value.entity_type == PiiEntityType.CREDIT_CARD assert excinfo.value.guardrail_name == presidio_guardrail.guardrail_name @@ -188,44 +185,46 @@ async def test_presidio_pre_call_hook_with_different_call_types(call_type): PiiEntityType.CREDIT_CARD: PiiAction.MASK, PiiEntityType.EMAIL_ADDRESS: PiiAction.MASK, } - + presidio_guardrail = _OPTIONAL_PresidioPIIMasking( pii_entities_config=pii_entities_config, presidio_analyzer_api_base=os.environ.get("PRESIDIO_ANALYZER_API_BASE"), - presidio_anonymizer_api_base=os.environ.get("PRESIDIO_ANONYMIZER_API_BASE") + presidio_anonymizer_api_base=os.environ.get("PRESIDIO_ANONYMIZER_API_BASE"), ) - + # Create a sample request with PII data data = { "messages": [ {"role": "system", "content": "You are a helpful assistant."}, - {"role": "user", "content": "My credit card is 4111-1111-1111-1111 and my email is test@example.com. My phone number is 555-123-4567"} + { + "role": "user", + "content": "My credit card is 4111-1111-1111-1111 and my email is test@example.com. My phone number is 555-123-4567", + }, ], - "model": "gpt-3.5-turbo" + "model": "gpt-3.5-turbo", } - + # Mock objects needed for the pre-call hook user_api_key_dict = UserAPIKeyAuth(api_key="test_key") cache = DualCache() - + # Call the pre-call hook with the specified call type modified_data = await presidio_guardrail.async_pre_call_hook( - user_api_key_dict=user_api_key_dict, - cache=cache, - data=data, - call_type=call_type + user_api_key_dict=user_api_key_dict, cache=cache, data=data, call_type=call_type ) - + # Verify the messages have been modified to mask PII - assert modified_data["messages"][0]["content"] == "You are a helpful assistant." # System prompt should be unchanged - + assert ( + modified_data["messages"][0]["content"] == "You are a helpful assistant." + ) # System prompt should be unchanged + user_message = modified_data["messages"][1]["content"] assert "4111-1111-1111-1111" not in user_message assert "test@example.com" not in user_message # Since this entity is not in the config, it should not be masked assert "555-123-4567" in user_message - + print(f"Modified user message for call_type={call_type}: {user_message}") @@ -243,7 +242,7 @@ def test_validate_environment_missing_http(base_url): # Use patch.dict to temporarily modify environment variables only for this test env_vars = { "PRESIDIO_ANALYZER_API_BASE": f"{base_url}/analyze", - "PRESIDIO_ANONYMIZER_API_BASE": f"{base_url}/anonymize" + "PRESIDIO_ANONYMIZER_API_BASE": f"{base_url}/anonymize", } with patch.dict(os.environ, env_vars): pii_masking.validate_environment() @@ -294,8 +293,12 @@ async def test_output_parsing(): new_response = await pii_masking.async_post_call_success_hook( user_api_key_dict=UserAPIKeyAuth(), data={ - "messages": [{"role": "system", "content": "You are an helpfull assistant"}], - "pii_tokens": {"": "Jane Doe", "": "034453334"}, + "messages": [ + {"role": "system", "content": "You are an helpfull assistant"} + ], + "metadata": { + "pii_tokens": {"": "Jane Doe", "": "034453334"} + }, }, response=response, ) @@ -440,24 +443,26 @@ async def test_presidio_pii_masking_logging_output_only_no_pre_api_hook(): @pytest.mark.asyncio -@patch.dict(os.environ, { - "PRESIDIO_ANALYZER_API_BASE": "http://localhost:5002", - "PRESIDIO_ANONYMIZER_API_BASE": "http://localhost:5001" -}) +@patch.dict( + os.environ, + { + "PRESIDIO_ANALYZER_API_BASE": "http://localhost:5002", + "PRESIDIO_ANONYMIZER_API_BASE": "http://localhost:5001", + }, +) async def test_presidio_pii_masking_logging_output_only_logged_response_guardrails_config(): from typing import Dict, List, Optional import litellm from litellm.proxy.guardrails.init_guardrails import initialize_guardrails from litellm.types.guardrails import ( - GuardrailItem, GuardrailItemSpec, GuardrailEventHooks, ) litellm.set_verbose = True # Environment variables are now patched via the decorator instead of setting them directly - + guardrails_config: List[Dict[str, GuardrailItemSpec]] = [ { "pii_masking": { @@ -499,60 +504,53 @@ async def test_presidio_pii_masking_logging_output_only_logged_response_guardrai async def test_presidio_language_configuration(): """Test that presidio_language parameter is properly set and used in analyze requests""" litellm._turn_on_debug() - + # Test with German language using mock testing to avoid API calls presidio_guardrail_de = _OPTIONAL_PresidioPIIMasking( pii_entities_config={}, presidio_language="de", - mock_testing=True # This bypasses the API validation + mock_testing=True, # This bypasses the API validation ) - + test_text = "Meine Telefonnummer ist +49 30 12345678" - + # Test the analyze request configuration analyze_request = presidio_guardrail_de._get_presidio_analyze_request_payload( - text=test_text, - presidio_config=None, - request_data={} + text=test_text, presidio_config=None, request_data={} ) - + # Verify the language is set to German assert analyze_request["language"] == "de" assert analyze_request["text"] == test_text - + # Test with Spanish language presidio_guardrail_es = _OPTIONAL_PresidioPIIMasking( - pii_entities_config={}, - presidio_language="es", - mock_testing=True + pii_entities_config={}, presidio_language="es", mock_testing=True ) - + test_text_es = "Mi número de teléfono es +34 912 345 678" - + analyze_request_es = presidio_guardrail_es._get_presidio_analyze_request_payload( - text=test_text_es, - presidio_config=None, - request_data={} + text=test_text_es, presidio_config=None, request_data={} ) - + # Verify the language is set to Spanish assert analyze_request_es["language"] == "es" assert analyze_request_es["text"] == test_text_es - + # Test default language (English) when not specified presidio_guardrail_default = _OPTIONAL_PresidioPIIMasking( - pii_entities_config={}, - mock_testing=True + pii_entities_config={}, mock_testing=True ) - + test_text_en = "My phone number is +1 555-123-4567" - - analyze_request_default = presidio_guardrail_default._get_presidio_analyze_request_payload( - text=test_text_en, - presidio_config=None, - request_data={} + + analyze_request_default = ( + presidio_guardrail_default._get_presidio_analyze_request_payload( + text=test_text_en, presidio_config=None, request_data={} + ) ) - + # Verify the language defaults to English assert analyze_request_default["language"] == "en" assert analyze_request_default["text"] == test_text_en @@ -562,36 +560,30 @@ async def test_presidio_language_configuration(): async def test_presidio_language_configuration_with_per_request_override(): """Test that per-request language configuration overrides the default configured language""" litellm._turn_on_debug() - + # Set up guardrail with German as default language presidio_guardrail = _OPTIONAL_PresidioPIIMasking( - pii_entities_config={}, - presidio_language="de", - mock_testing=True + pii_entities_config={}, presidio_language="de", mock_testing=True ) - + test_text = "Test text with PII" - + # Test with per-request config overriding the default language presidio_config = PresidioPerRequestConfig(language="fr") - + analyze_request = presidio_guardrail._get_presidio_analyze_request_payload( - text=test_text, - presidio_config=presidio_config, - request_data={} + text=test_text, presidio_config=presidio_config, request_data={} ) - + # Verify the per-request language (French) overrides the default (German) assert analyze_request["language"] == "fr" assert analyze_request["text"] == test_text - + # Test without per-request config - should use default language analyze_request_default = presidio_guardrail._get_presidio_analyze_request_payload( - text=test_text, - presidio_config=None, - request_data={} + text=test_text, presidio_config=None, request_data={} ) - + # Verify the default language (German) is used assert analyze_request_default["language"] == "de" assert analyze_request_default["text"] == test_text 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..32a8c1b1070 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,625 @@ 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) + 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"] + seq = len(pii_tokens) + 1 + token = f"" + 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 = "" + 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?" + + +@pytest.mark.parametrize( + "initial_hook", + ["pre_call", "during_call", "pre_mcp_call"], +) +def test_event_hook_auto_expansion_for_all_string_hooks(initial_hook): + """ + Regression test: when output_parse_pii is True, the guardrail must add + 'post_call' to event_hook regardless of the initial string hook value, + not just when it's 'pre_call'. + """ + guardrail = _OPTIONAL_PresidioPIIMasking( + mock_testing=True, + output_parse_pii=True, + event_hook=initial_hook, + ) + assert isinstance(guardrail.event_hook, list) + assert initial_hook in guardrail.event_hook + assert "post_call" in guardrail.event_hook + + +def test_event_hook_no_expansion_when_already_post_call(): + """post_call alone should stay as-is — no expansion needed.""" + guardrail = _OPTIONAL_PresidioPIIMasking( + mock_testing=True, + output_parse_pii=True, + event_hook="post_call", + ) + # Should remain a string "post_call", not expanded to a list + assert guardrail.event_hook == "post_call" + + +@pytest.mark.asyncio +async def test_metadata_none_does_not_crash(): + """ + Regression test: if metadata is explicitly None in request_data, + the guardrail must not crash with TypeError on the write or read path. + """ + guardrail = _OPTIONAL_PresidioPIIMasking( + mock_testing=True, + output_parse_pii=True, + ) + + token_key = "" + # metadata explicitly None — must not crash + request_data = { + "model": "gpt-3.5-turbo", + "metadata": None, + } + + response = ModelResponse( + choices=[ + Choices( + message=Message( + role="assistant", + content=f"Hello {token_key}, how can I help you?", + ), + index=0, + finish_reason="stop", + ) + ] + ) + + # Should not raise TypeError + await guardrail._process_response_for_pii( + response=response, + request_data=request_data, + mode="unmask", + ) + + # No pii_tokens to unmask, so content stays as-is + assert ( + response.choices[0].message.content == f"Hello {token_key}, how can I help you?" + ) + + +# --------------------------------------------------------------------------- +# Tests for sequential-numbered token unmasking in _unmask_pii_text +# --------------------------------------------------------------------------- + + +def test_unmask_exact_match_with_sequential_tokens(): + """ + Normal unmasking: LLM echoes numbered tokens verbatim → original PII restored. + """ + from litellm.proxy.guardrails.guardrail_hooks.presidio import ( + _OPTIONAL_PresidioPIIMasking, + ) + + pii_tokens = { + "": "John Smith", + "": "555-123-4567", + } + text = "Hello , your number is ." + result = _OPTIONAL_PresidioPIIMasking._unmask_pii_text(text, pii_tokens) + assert result == "Hello John Smith, your number is 555-123-4567." + + +def test_unmask_multiple_same_entity_type(): + """ + Two phone numbers get distinct numbered tokens and unmask correctly. + """ + from litellm.proxy.guardrails.guardrail_hooks.presidio import ( + _OPTIONAL_PresidioPIIMasking, + ) + + pii_tokens = { + "": "555-111-0000", + "": "555-222-0000", + } + text = "Call or ." + result = _OPTIONAL_PresidioPIIMasking._unmask_pii_text(text, pii_tokens) + assert result == "Call 555-111-0000 or 555-222-0000." + + +def test_unmask_graceful_degradation(): + """ + If the LLM doesn't echo the token back, the numbered label stays + in the output — clean and readable, not garbage hex. + """ + from litellm.proxy.guardrails.guardrail_hooks.presidio import ( + _OPTIONAL_PresidioPIIMasking, + ) + + pii_tokens = { + "": "John", + } + # LLM paraphrased instead of echoing the token + text = "I see you provided a name." + result = _OPTIONAL_PresidioPIIMasking._unmask_pii_text(text, pii_tokens) + # No change — no garbage, just clean text + assert result == text + + +# --------------------------------------------------------------------------- +# Fix 1: Position bug — reverse sort + original text coordinates +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_anonymize_text_multiple_items_position_correctness(): + """ + Regression test: when multiple PII items exist, coordinates reference the + ORIGINAL text. Processing in reverse order prevents coordinate drift. + """ + guardrail = _OPTIONAL_PresidioPIIMasking( + presidio_analyzer_api_base="http://test-analyzer/", + presidio_anonymizer_api_base="http://test-anonymizer/", + mock_testing=False, + ) + + # "Call John at 555-123-4567" + # "John" at [5:9], "555-123-4567" at [13:25] + anonymizer_response = { + "text": "Call at ", + "items": [ + { + "start": 5, + "end": 9, + "entity_type": "PERSON", + "text": "", + "operator": "replace", + }, + { + "start": 13, + "end": 25, + "entity_type": "PHONE_NUMBER", + "text": "", + "operator": "replace", + }, + ], + } + + mock_iterator = _make_mock_session_iterator(anonymizer_response) + + request_data = {"metadata": {}} + with patch.object(guardrail, "_get_session_iterator", mock_iterator): + result = await guardrail.anonymize_text( + text="Call John at 555-123-4567", + analyze_results=[ + {"start": 5, "end": 9, "entity_type": "PERSON", "score": 0.9}, + {"start": 13, "end": 25, "entity_type": "PHONE_NUMBER", "score": 0.95}, + ], + output_parse_pii=True, + masked_entity_count={}, + request_data=request_data, + ) + + pii_tokens = request_data["metadata"]["pii_tokens"] + + # Verify tokens captured the correct ORIGINAL text values + person_token = [k for k in pii_tokens if "PERSON" in k][0] + phone_token = [k for k in pii_tokens if "PHONE" in k][0] + assert pii_tokens[person_token] == "John" + assert pii_tokens[phone_token] == "555-123-4567" + + # Verify both PII values are masked in the result + assert "John" not in result + assert "555-123-4567" not in result + + +# --------------------------------------------------------------------------- +# Fix 2: Anthropic native dict response handling +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_anthropic_native_response_unmasking(): + """ + Anthropic native dict responses (type='message') should be unmasked + when output_parse_pii is enabled. + """ + guardrail = _OPTIONAL_PresidioPIIMasking( + mock_testing=True, + output_parse_pii=True, + ) + + request_data = { + "model": "claude-3-haiku", + "metadata": { + "pii_tokens": { + "": "John Smith", + "": "555-123-4567", + } + }, + } + + anthropic_response = { + "type": "message", + "id": "msg_123", + "model": "claude-3-haiku", + "role": "assistant", + "content": [ + { + "type": "text", + "text": "Hello , your number is .", + } + ], + "stop_reason": "end_turn", + "usage": {"input_tokens": 10, "output_tokens": 20}, + } + + mock_user_api_key = UserAPIKeyAuth(api_key="test-key") + + result = await guardrail.async_post_call_success_hook( + data=request_data, + user_api_key_dict=mock_user_api_key, + response=anthropic_response, + ) + + assert result["content"][0]["text"] == ( + "Hello John Smith, your number is 555-123-4567." + ) + + +@pytest.mark.asyncio +async def test_anthropic_native_response_masking(): + """ + Anthropic native dict responses should be masked when + apply_to_output is enabled. + """ + guardrail = _OPTIONAL_PresidioPIIMasking( + mock_testing=True, + apply_to_output=True, + ) + + async def mock_check_pii(text, output_parse_pii, presidio_config, request_data): + return text.replace("John Smith", "[PERSON]").replace("555-123-4567", "[PHONE]") + + guardrail.check_pii = mock_check_pii + + anthropic_response = { + "type": "message", + "id": "msg_123", + "model": "claude-3-haiku", + "role": "assistant", + "content": [{"type": "text", "text": "Hello John Smith, call 555-123-4567."}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 10, "output_tokens": 20}, + } + + mock_user_api_key = UserAPIKeyAuth(api_key="test-key") + + result = await guardrail.async_post_call_success_hook( + data={}, + user_api_key_dict=mock_user_api_key, + response=anthropic_response, + ) + + assert "[PERSON]" in result["content"][0]["text"] + assert "[PHONE]" in result["content"][0]["text"] + assert "John Smith" not in result["content"][0]["text"] + + +@pytest.mark.asyncio +async def test_anthropic_native_response_non_text_blocks_untouched(): + """ + Non-text blocks (tool_use, thinking) in Anthropic responses + should be left untouched during unmasking. + """ + guardrail = _OPTIONAL_PresidioPIIMasking( + mock_testing=True, + output_parse_pii=True, + ) + + request_data = { + "model": "claude-3-haiku", + "metadata": {"pii_tokens": {"": "John"}}, + } + + anthropic_response = { + "type": "message", + "id": "msg_123", + "content": [ + {"type": "text", "text": "Hello "}, + { + "type": "tool_use", + "id": "call_1", + "name": "search", + "input": {"q": "test"}, + }, + ], + "role": "assistant", + "stop_reason": "end_turn", + "usage": {"input_tokens": 10, "output_tokens": 20}, + } + + mock_user_api_key = UserAPIKeyAuth(api_key="test-key") + + result = await guardrail.async_post_call_success_hook( + data=request_data, + user_api_key_dict=mock_user_api_key, + response=anthropic_response, + ) + + assert result["content"][0]["text"] == "Hello John" + assert result["content"][1]["type"] == "tool_use" + assert result["content"][1]["name"] == "search" + + +# --------------------------------------------------------------------------- +# Fix 3: Anthropic native SSE streaming — bytes passthrough +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_streaming_bytes_chunks_are_yielded_not_discarded(): + """ + Regression test: bytes chunks (Anthropic native SSE) should be yielded + through the streaming hook, not silently discarded. + """ + + guardrail = _OPTIONAL_PresidioPIIMasking( + mock_testing=True, + apply_to_output=True, + ) + + byte_chunk = b'data: {"type":"content_block_delta","delta":{"text":"Hello"}}\n\n' + + async def mock_stream(): + yield byte_chunk + + mock_user_api_key = UserAPIKeyAuth(api_key="test-key") + chunks = [] + async for chunk in guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=mock_user_api_key, + response=mock_stream(), + request_data={}, + ): + chunks.append(chunk) + + assert any( + isinstance(c, bytes) for c in chunks + ), "bytes chunks must not be discarded" + assert byte_chunk in chunks + + +@pytest.mark.asyncio +async def test_streaming_unmask_path_bytes_passthrough(): + """ + Bytes chunks in the unmasking path should also pass through. + """ + guardrail = _OPTIONAL_PresidioPIIMasking( + mock_testing=True, + output_parse_pii=True, + ) + + byte_chunk = b'data: {"type":"content_block_delta"}\n\n' + request_data = { + "metadata": {"pii_tokens": {"": "John"}}, + } + + async def mock_stream(): + yield byte_chunk + + mock_user_api_key = UserAPIKeyAuth(api_key="test-key") + chunks = [] + async for chunk in guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=mock_user_api_key, + response=mock_stream(), + request_data=request_data, + ): + chunks.append(chunk) + + assert len(chunks) == 1 + assert chunks[0] == byte_chunk + + +# --------------------------------------------------------------------------- +# Fix 4: apply_guardrail unmask path for input_type="response" +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_apply_guardrail_unmask_on_response(): + """ + When input_type is 'response' and pii_tokens exist, apply_guardrail + should unmask text instead of masking it. + """ + guardrail = _OPTIONAL_PresidioPIIMasking( + guardrail_name="test_presidio", + output_parse_pii=True, + mock_testing=True, + ) + + request_data = { + "model": "gpt-4o", + "metadata": { + "pii_tokens": { + "": "John Smith", + "": "555-123-4567", + } + }, + } + + inputs = { + "texts": [ + "Hello , your number is .", + ] + } + + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data=request_data, + input_type="response", + ) + + assert result["texts"][0] == "Hello John Smith, your number is 555-123-4567." + + +@pytest.mark.asyncio +async def test_apply_guardrail_masks_on_request(): + """ + When input_type is 'request', apply_guardrail should mask as before. + """ + guardrail = _OPTIONAL_PresidioPIIMasking( + guardrail_name="test_presidio", + output_parse_pii=True, + mock_testing=True, + ) + + async def mock_check_pii(text, output_parse_pii, presidio_config, request_data): + return text.replace("John Smith", "") + + guardrail.check_pii = mock_check_pii + + result = await guardrail.apply_guardrail( + inputs={"texts": ["Hello John Smith"]}, + request_data={"model": "gpt-4o", "metadata": {}}, + input_type="request", + ) + + assert "" in result["texts"][0] + assert "John Smith" not in result["texts"][0] + + +@pytest.mark.asyncio +async def test_apply_to_output_streaming_bytes_only_logs_warning(): + """ + Regression test: when apply_to_output=True and the stream contains only + bytes chunks (Anthropic native SSE), output masking is skipped. + A warning must be logged so operators are aware. + """ + guardrail = _OPTIONAL_PresidioPIIMasking( + mock_testing=True, + apply_to_output=True, + ) + + byte_chunks = [ + b'data: {"type":"content_block_delta","delta":{"text":"Hello"}}\n\n', + b'data: {"type":"content_block_delta","delta":{"text":" world"}}\n\n', + ] + + async def mock_stream(): + for b in byte_chunks: + yield b + + mock_user_api_key = UserAPIKeyAuth(api_key="test-key") + + collected = [] + with patch( + "litellm.proxy.guardrails.guardrail_hooks.presidio.verbose_proxy_logger" + ) as mock_logger: + async for chunk in guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=mock_user_api_key, + response=mock_stream(), + request_data={}, + ): + collected.append(chunk) + + # All bytes should be yielded through + assert len(collected) == len(byte_chunks) + for original, received in zip(byte_chunks, collected): + assert original == received + + # Warning must be logged about skipped masking + mock_logger.warning.assert_called_once() + warning_msg = mock_logger.warning.call_args[0][0] + assert "Output PII masking was skipped" in warning_msg