mirror of
https://github.com/tiennm99/litellm.git
synced 2026-08-05 12:23:28 +00:00
fix test case and req changes
This commit is contained in:
@@ -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):
|
||||
"""
|
||||
|
||||
Reference in New Issue
Block a user