ai-workspace-services/litellm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: presidio streaming, false positives

Browse Source
This commit is contained in:
Harshit28j 2026-02-24 00:42:29 +05:30
parent 4dfec4661b
commit af9ad68a43
4 changed files with 294 additions and 109 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

25
docs/my-website/docs/tutorials/presidio_pii_masking.md
View File

@ -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
```

205
litellm/proxy/guardrails/guardrail_hooks/presidio.py
View File

@ -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(

8
litellm/types/guardrails.py
View File

@ -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",

165
tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py
View File

@ -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"