From e5c907dc93f32066fa28b9d0845feffca2763634 Mon Sep 17 00:00:00 2001
From: Harshit28j <Harshitjain0562@gmail.com>
Date: Tue, 24 Feb 2026 09:13:27 +0530
Subject: [PATCH] fix test case and req changes

---
 .../guardrails/guardrail_hooks/presidio.py    | 16 ++--
 .../guardrail_hooks/test_presidio.py          | 81 -------------------
 2 files changed, 10 insertions(+), 87 deletions(-)

diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py
index dbfa2b51de..28ccff0e36 100644
--- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py
+++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py
@@ -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
+            )
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 43bc47834e..e2c05f0ad3 100644
--- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py
+++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py
@@ -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):
     """