From af9ad68a43228ff368196d2e10198bd79d650a63 Mon Sep 17 00:00:00 2001 From: Harshit28j Date: Tue, 24 Feb 2026 00:42:29 +0530 Subject: [PATCH] fix: presidio streaming, false positives --- .../docs/tutorials/presidio_pii_masking.md | 25 +++ .../guardrails/guardrail_hooks/presidio.py | 205 ++++++++---------- litellm/types/guardrails.py | 8 + .../guardrail_hooks/test_presidio.py | 165 ++++++++++++++ 4 files changed, 294 insertions(+), 109 deletions(-) diff --git a/docs/my-website/docs/tutorials/presidio_pii_masking.md b/docs/my-website/docs/tutorials/presidio_pii_masking.md index 315639d8d66..ea3761163f2 100644 --- a/docs/my-website/docs/tutorials/presidio_pii_masking.md +++ b/docs/my-website/docs/tutorials/presidio_pii_masking.md @@ -685,3 +685,28 @@ Congratulations! 🎉 You've successfully set up PII masking with Presidio and L --- **Need help?** Join our [Discord community](https://discord.com/invite/wuPM9dRgDw) or open an issue on GitHub! + +### Suppressing False Positives + +Presidio can sometimes trigger false positive detections. For example, short alphanumeric strings might be incorrectly flagged as `US_DRIVER_LICENSE`. + +You can suppress these false positives using `presidio_score_thresholds` or `presidio_entities_deny_list`. + +```yaml +guardrails: + - guardrail_name: presidio-pii + litellm_params: + guardrail: presidio + mode: "pre_call" + presidio_analyzer_api_base: "http://localhost:5002/" + presidio_anonymizer_api_base: "http://localhost:5001/" + + # Use high score thresholds to reduce false positives + presidio_score_thresholds: + US_DRIVER_LICENSE: 0.85 + ALL: 0.5 + + # Or exclude certain entity types entirely from detection + presidio_entities_deny_list: + - US_DRIVER_LICENSE +``` diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index 29d57ee4734..dbfa2b51de3 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -87,6 +87,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): presidio_score_thresholds: Optional[ Dict[Union[PiiEntityType, str], float] ] = None, + presidio_entities_deny_list: Optional[List[Union[PiiEntityType, str]]] = None, **kwargs, ): if logging_only is True: @@ -106,6 +107,9 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): self.presidio_score_thresholds: Dict[Union[PiiEntityType, str], float] = ( presidio_score_thresholds or {} ) + self.presidio_entities_deny_list: List[Union[PiiEntityType, str]] = ( + presidio_entities_deny_list or [] + ) self.presidio_language = presidio_language or "en" # Shared HTTP session to prevent memory leaks (issue #14540) self._http_session: Optional[aiohttp.ClientSession] = None @@ -391,9 +395,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): except Exception as e: # Sanitize exception to avoid leaking the original text (which may # contain API keys or other secrets) in error responses. - raise Exception( - f"Presidio PII analysis failed: {type(e).__name__}" - ) from e + raise Exception(f"Presidio PII analysis failed: {type(e).__name__}") from e async def anonymize_text( self, @@ -464,9 +466,10 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): self, analyze_results: Union[List[PresidioAnalyzeResponseItem], Dict] ) -> Union[List[PresidioAnalyzeResponseItem], Dict]: """ - Drop detections that fall below configured per-entity score thresholds. + Drop detections that fall below configured per-entity score thresholds + or match an entity type in the deny list. """ - if not self.presidio_score_thresholds: + if not self.presidio_score_thresholds and not self.presidio_entities_deny_list: return analyze_results if not isinstance(analyze_results, list): @@ -475,17 +478,21 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): filtered_results: List[PresidioAnalyzeResponseItem] = [] for item in analyze_results: entity_type = item.get("entity_type") - score = item.get("score") - threshold = None - if entity_type is not None: - threshold = self.presidio_score_thresholds.get(entity_type) - if threshold is None: - threshold = self.presidio_score_thresholds.get("ALL") + if entity_type and entity_type in self.presidio_entities_deny_list: + continue - if threshold is not None: - if score is None or score < threshold: - continue + if self.presidio_score_thresholds: + score = item.get("score") + threshold = None + if entity_type is not None: + threshold = self.presidio_score_thresholds.get(entity_type) + if threshold is None: + threshold = self.presidio_score_thresholds.get("ALL") + + if threshold is not None: + if score is None or score < threshold: + continue filtered_results.append(item) @@ -619,9 +626,9 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): if messages is None: return data tasks = [] - task_mappings: List[Tuple[int, Optional[int]]] = ( - [] - ) # Track (message_index, content_index) for each task + task_mappings: List[ + Tuple[int, Optional[int]] + ] = [] # Track (message_index, content_index) for each task for msg_idx, m in enumerate(messages): content = m.get("content", None) @@ -722,9 +729,9 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): ): # /chat/completions requests messages: Optional[List] = kwargs.get("messages", None) tasks = [] - task_mappings: List[Tuple[int, Optional[int]]] = ( - [] - ) # Track (message_index, content_index) for each task + task_mappings: List[ + Tuple[int, Optional[int]] + ] = [] # Track (message_index, content_index) for each task if messages is None: return kwargs, result @@ -877,66 +884,69 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): ) -> AsyncGenerator[ModelResponseStream, None]: """ Process streaming response chunks to unmask PII tokens when needed. - - If PII processing is enabled, this collects all chunks, applies PII unmasking, - and returns a reconstructed stream. Otherwise, it passes through the original stream. """ - # If we need to mask model output, collect the full stream, apply masking, and replay it. + from litellm.llms.base_llm.base_model_iterator import MockResponseIterator + from litellm.main import stream_chunk_builder + from litellm.types.utils import ModelResponse + + # --- Output masking path (apply_to_output=True) --- if self.apply_to_output: - from litellm.llms.base_llm.base_model_iterator import MockResponseIterator - from litellm.types.utils import Choices, Message - try: - collected_content = "" - last_chunk = None - + all_chunks: List[ModelResponseStream] = [] async for chunk in response: - last_chunk = chunk + if isinstance(chunk, ModelResponseStream): + all_chunks.append(chunk) - if ( - hasattr(chunk, "choices") - and chunk.choices - and hasattr(chunk.choices[0], "delta") - and hasattr(chunk.choices[0].delta, "content") - and isinstance(chunk.choices[0].delta.content, str) - ): - collected_content += chunk.choices[0].delta.content + if not all_chunks: + return - if not last_chunk: - async for chunk in response: + assembled_model_response = stream_chunk_builder( + chunks=all_chunks, messages=request_data.get("messages") + ) + + if not isinstance(assembled_model_response, ModelResponse): + for chunk in all_chunks: yield chunk return + # Apply Presidio masking on the assembled response presidio_config = self.get_presidio_settings_from_request_data( request_data or {} ) + + content_to_mask = "" + if ( + hasattr(assembled_model_response, "choices") + and len(assembled_model_response.choices) > 0 + ): + if hasattr( + assembled_model_response.choices[0], "message" + ) and hasattr( + assembled_model_response.choices[0].message, "content" + ): + content_to_mask = ( + assembled_model_response.choices[0].message.content or "" + ) + masked_content = await self.check_pii( - text=collected_content, + text=content_to_mask, output_parse_pii=False, presidio_config=presidio_config, request_data=request_data, ) - mock_response = MockResponseIterator( - model_response=ModelResponse( - id=last_chunk.id, - object=last_chunk.object, - created=last_chunk.created, - model=last_chunk.model, - choices=[ - Choices( - message=Message( - role="assistant", - content=masked_content, - ), - index=0, - finish_reason="stop", - ) - ], - ), - json_mode=False, - ) + if ( + hasattr(assembled_model_response, "choices") + and len(assembled_model_response.choices) > 0 + ): + if hasattr(assembled_model_response.choices[0], "message"): + assembled_model_response.choices[ + 0 + ].message.content = masked_content + mock_response = MockResponseIterator( + model_response=assembled_model_response + ) async for chunk in mock_response: yield chunk return @@ -945,77 +955,54 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): verbose_proxy_logger.error( f"Error masking streaming PII output: {str(e)}" ) - async for chunk in response: + # Cannot re-iterate `response` — it's already consumed. + # If we collected chunks before the error, replay those. + for chunk in all_chunks: yield chunk return - # If PII unmasking not needed, just pass through the original stream + # --- PII unmasking path (output_parse_pii=True) --- if not (self.output_parse_pii and self.pii_tokens): async for chunk in response: yield chunk return - # Import here to avoid circular imports - from litellm.llms.base_llm.base_model_iterator import MockResponseIterator - from litellm.types.utils import Choices, Message - try: - # Collect all chunks to process them together - collected_content = "" - last_chunk = None - + all_chunks: List[ModelResponseStream] = [] async for chunk in response: - last_chunk = chunk + if isinstance(chunk, ModelResponseStream): + all_chunks.append(chunk) - # Extract content safely with proper attribute checks - if ( - hasattr(chunk, "choices") - and chunk.choices - and hasattr(chunk.choices[0], "delta") - and hasattr(chunk.choices[0].delta, "content") - and isinstance(chunk.choices[0].delta.content, str) - ): - collected_content += chunk.choices[0].delta.content + if not all_chunks: + return - # No need to proceed if we didn't capture a valid chunk - if not last_chunk: - async for chunk in response: + assembled_model_response = stream_chunk_builder( + chunks=all_chunks, messages=request_data.get("messages") + ) + + if not isinstance(assembled_model_response, ModelResponse): + for chunk in all_chunks: yield chunk return - # Apply PII unmasking to the complete content - for token, original_text in self.pii_tokens.items(): - collected_content = collected_content.replace(token, original_text) + # Apply PII unmasking to assembled content + for choice in assembled_model_response.choices: + if hasattr(choice, "message") and hasattr(choice.message, "content"): + content = choice.message.content + if isinstance(content, str): + for token, original_text in self.pii_tokens.items(): + content = content.replace(token, original_text) + choice.message.content = content - # Reconstruct the response with unmasked content mock_response = MockResponseIterator( - model_response=ModelResponse( - id=last_chunk.id, - object=last_chunk.object, - created=last_chunk.created, - model=last_chunk.model, - choices=[ - Choices( - message=Message( - role="assistant", - content=collected_content, - ), - index=0, - finish_reason="stop", - ) - ], - ), - json_mode=False, + model_response=assembled_model_response ) - - # Return the reconstructed stream async for chunk in mock_response: yield chunk except Exception as e: verbose_proxy_logger.error(f"Error in PII streaming processing: {str(e)}") - # Fallback to original stream on error - async for chunk in response: + for chunk in all_chunks: yield chunk def get_presidio_settings_from_request_data( diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index c87d1076d0e..c14bb62fa28 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -309,6 +309,14 @@ class PresidioConfigModel(PresidioPresidioConfigModelUserInterface): "Entities below the threshold are ignored." ), ) + presidio_entities_deny_list: Optional[List[Union[PiiEntityType, str]]] = Field( + default=None, + description=( + "List of entity types to exclude from Presidio detection results. " + "Detections of these types will be silently dropped. " + "Useful for suppressing false positives (e.g., US_DRIVER_LICENSE on coding routes)." + ), + ) presidio_ad_hoc_recognizers: Optional[str] = Field( default=None, description="Path to a JSON file containing ad-hoc recognizers for Presidio", 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 f01c23f7116..43bc47834e0 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py @@ -1360,3 +1360,168 @@ async def test_get_session_iterator_thread_safety(presidio_guardrail): assert not bg_session.closed, "Background session should remain open for reuse" print("✓ Session iterator thread safety test passed") + + +from litellm.types.utils import ModelResponseStream + + +@pytest.mark.asyncio +async def test_streaming_with_bytes_chunks_does_not_crash(mock_user_api_key): + """ + Regression test: async_post_call_streaming_iterator_hook should + gracefully handle raw bytes in the stream instead of crashing with + 'bytes' object has no attribute 'id'. + """ + guardrail = _OPTIONAL_PresidioPIIMasking( + mock_testing=True, + apply_to_output=True, + mock_redacted_text={"text": "redacted"}, + ) + + async def mock_stream(): + yield b'data: {"id":"chatcmpl-1"}\n\n' # raw bytes + yield ModelResponseStream( + id="chatcmpl-1", + choices=[], + created=1, + model="gpt-4", + object="chat.completion.chunk", + system_fingerprint=None, + ) # proper chunk + + 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) + + # Should not crash, should produce at least one valid chunk + assert len(chunks) >= 1 + + +def test_entity_deny_list_filters_detections(): + """ + Verify presidio_entities_deny_list removes matching entity types. + """ + guardrail = _OPTIONAL_PresidioPIIMasking( + mock_testing=True, + presidio_entities_deny_list=["US_DRIVER_LICENSE"], + ) + + results = [ + {"entity_type": "US_DRIVER_LICENSE", "start": 0, "end": 2, "score": 0.6}, + {"entity_type": "CREDIT_CARD", "start": 10, "end": 26, "score": 0.95}, + ] + + filtered = guardrail.filter_analyze_results_by_score(results) + + assert len(filtered) == 1 + assert filtered[0]["entity_type"] == "CREDIT_CARD" + + +def test_deny_list_and_score_threshold_combined(): + """ + Verify deny list + score threshold work together correctly. + """ + guardrail = _OPTIONAL_PresidioPIIMasking( + mock_testing=True, + presidio_entities_deny_list=["US_DRIVER_LICENSE"], + presidio_score_thresholds={"ALL": 0.8}, + ) + + results = [ + {"entity_type": "US_DRIVER_LICENSE", "start": 0, "end": 2, "score": 0.95}, + {"entity_type": "CREDIT_CARD", "start": 10, "end": 26, "score": 0.6}, + {"entity_type": "EMAIL_ADDRESS", "start": 30, "end": 50, "score": 0.9}, + ] + + filtered = guardrail.filter_analyze_results_by_score(results) + + # US_DRIVER_LICENSE excluded by deny list (even though score > 0.8) + # CREDIT_CARD excluded by score threshold (0.6 < 0.8) + # EMAIL_ADDRESS passes both filters + assert len(filtered) == 1 + assert filtered[0]["entity_type"] == "EMAIL_ADDRESS" + + +@pytest.mark.asyncio +async def test_streaming_with_bytes_chunks_does_not_crash(mock_user_api_key): + """ + Regression test: async_post_call_streaming_iterator_hook should + gracefully handle raw bytes in the stream instead of crashing with + 'bytes' object has no attribute 'id'. + """ + guardrail = _OPTIONAL_PresidioPIIMasking( + mock_testing=True, + apply_to_output=True, + mock_redacted_text={"text": "redacted"}, + ) + + async def mock_stream(): + yield b'data: {"id":"chatcmpl-1"}\n\n' # raw bytes + yield ModelResponseStream( + id="chatcmpl-1", + choices=[], + created=1, + model="gpt-4", + object="chat.completion.chunk", + system_fingerprint=None, + ) # proper chunk + + 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) + + # Should not crash, should produce at least one valid chunk + assert len(chunks) >= 1 + + +def test_entity_deny_list_filters_detections(): + """ + Verify presidio_entities_deny_list removes matching entity types. + """ + guardrail = _OPTIONAL_PresidioPIIMasking( + mock_testing=True, + presidio_entities_deny_list=["US_DRIVER_LICENSE"], + ) + + results = [ + {"entity_type": "US_DRIVER_LICENSE", "start": 0, "end": 2, "score": 0.6}, + {"entity_type": "CREDIT_CARD", "start": 10, "end": 26, "score": 0.95}, + ] + + filtered = guardrail.filter_analyze_results_by_score(results) + + assert len(filtered) == 1 + assert filtered[0]["entity_type"] == "CREDIT_CARD" + + +def test_deny_list_and_score_threshold_combined(): + """ + Verify deny list + score threshold work together correctly. + """ + guardrail = _OPTIONAL_PresidioPIIMasking( + mock_testing=True, + presidio_entities_deny_list=["US_DRIVER_LICENSE"], + presidio_score_thresholds={"ALL": 0.8}, + ) + + results = [ + {"entity_type": "US_DRIVER_LICENSE", "start": 0, "end": 2, "score": 0.95}, + {"entity_type": "CREDIT_CARD", "start": 10, "end": 26, "score": 0.6}, + {"entity_type": "EMAIL_ADDRESS", "start": 30, "end": 50, "score": 0.9}, + ] + + filtered = guardrail.filter_analyze_results_by_score(results) + + # US_DRIVER_LICENSE excluded by deny list (even though score > 0.8) + # CREDIT_CARD excluded by score threshold (0.6 < 0.8) + # EMAIL_ADDRESS passes both filters + assert len(filtered) == 1 + assert filtered[0]["entity_type"] == "EMAIL_ADDRESS"