fix test case and req changes

This commit is contained in:
Harshit28j
2026-02-24 09:13:27 +05:30
parent af9ad68a43
commit e5c907dc93
2 changed files with 10 additions and 87 deletions
@@ -968,20 +968,20 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
return
try:
all_chunks: List[ModelResponseStream] = []
remaining_chunks: List[ModelResponseStream] = []
async for chunk in response:
if isinstance(chunk, ModelResponseStream):
all_chunks.append(chunk)
remaining_chunks.append(chunk)
if not all_chunks:
if not remaining_chunks:
return
assembled_model_response = stream_chunk_builder(
chunks=all_chunks, messages=request_data.get("messages")
chunks=remaining_chunks, messages=request_data.get("messages")
)
if not isinstance(assembled_model_response, ModelResponse):
for chunk in all_chunks:
for chunk in remaining_chunks:
yield chunk
return
@@ -1002,7 +1002,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
except Exception as e:
verbose_proxy_logger.error(f"Error in PII streaming processing: {str(e)}")
for chunk in all_chunks:
for chunk in remaining_chunks:
yield chunk
def get_presidio_settings_from_request_data(
@@ -1063,3 +1063,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
self.pii_entities_config = litellm_params.pii_entities_config
if litellm_params.presidio_score_thresholds:
self.presidio_score_thresholds = litellm_params.presidio_score_thresholds
if litellm_params.presidio_entities_deny_list:
self.presidio_entities_deny_list = (
litellm_params.presidio_entities_deny_list
)
@@ -1365,87 +1365,6 @@ async def test_get_session_iterator_thread_safety(presidio_guardrail):
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):
"""