fix(presidio): preserve span union and prevent PII exposure in overlap resolution

litellm/proxy/guardrails/guardrail_hooks/presidio.py

@@ -9,11 +9,13 @@
 
 
 import asyncio
-import json
-import threading
-from collections.abc import AsyncGenerator, AsyncIterable, Awaitable, Sequence
+from collections.abc import AsyncGenerator, AsyncIterable, Awaitable, Mapping, Sequence
 from contextlib import asynccontextmanager
 from datetime import datetime
+from functools import reduce
+import json
+import threading
+from types import MappingProxyType
 from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, TypedDict, cast
 
 import aiohttp
@@ -695,32 +697,57 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
 
     @staticmethod
     def _resolve_overlapping_spans(
-        analyze_results: list[dict[str, Any]] | list[PresidioAnalyzeResponseItem],
-    ) -> list[dict[str, Any]]:
-        if not analyze_results:
-            return []
-
-        sorted_by_priority = sorted(
-            analyze_results,
-            key=lambda x: (
-                -x.get("score", 0),
-                -(x.get("end", 0) - x.get("start", 0)),
-                x.get("start", 0),
-            ),
+        analyze_results: Sequence[Mapping[str, Any]],
+    ) -> tuple[Mapping[str, Any], ...]:
+        valid_candidates: Final = tuple(
+            c
+            for c in analyze_results
+            if c.get("start") is not None
+            and c.get("end") is not None
+            and int(c["start"]) < int(c["end"])
         )
-        accepted_spans: list[dict[str, Any]] = []
-        for candidate in sorted_by_priority:
-            c_start = candidate["start"]
-            c_end = candidate["end"]
-            overlaps = False
-            for acc in accepted_spans:
-                if max(c_start, acc["start"]) < min(c_end, acc["end"]):
-                    overlaps = True
-                    break
-            if not overlaps:
-                accepted_spans.append(candidate)
+        if not valid_candidates:
+            return ()
 
-        return sorted(accepted_spans, key=lambda x: x["start"])
+        sorted_by_start: Final = sorted(
+            valid_candidates,
+            key=lambda x: (int(x["start"]), int(x["end"])),
+        )
+
+        def _cluster_reducer(
+            acc: tuple[tuple[Mapping[str, Any], ...], ...],
+            item: Mapping[str, Any],
+        ) -> tuple[tuple[Mapping[str, Any], ...], ...]:
+            if not acc:
+                return ((item,),)
+            last_cluster: Final = acc[-1]
+            cluster_max_end: Final = max(int(s["end"]) for s in last_cluster)
+            if int(item["start"]) < cluster_max_end:
+                return acc[:-1] + (last_cluster + (item,),)
+            return acc + ((item,),)
+
+        clusters: Final = reduce(_cluster_reducer, sorted_by_start, ())
+
+        def _resolve_cluster(cluster: tuple[Mapping[str, Any], ...]) -> Mapping[str, Any]:
+            cluster_start: Final = min(int(s["start"]) for s in cluster)
+            cluster_end: Final = max(int(s["end"]) for s in cluster)
+            best_candidate: Final = max(
+                cluster,
+                key=lambda x: (
+                    int(x["start"]) == cluster_start and int(x["end"]) == cluster_end,
+                    float(x.get("score") or 0),
+                    int(x["end"]) - int(x["start"]),
+                ),
+            )
+            resolved: Final = MappingProxyType({
+                "entity_type": best_candidate.get("entity_type") or "UNKNOWN",
+                "start": cluster_start,
+                "end": cluster_end,
+                "score": max(float(s.get("score") or 0) for s in cluster),
+            })
+            return resolved
+
+        return tuple(_resolve_cluster(c) for c in clusters)
 
     def _finalize_presidio_anonymize_numbered_tokens(
         self,

tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py

@@ -3196,15 +3196,45 @@ def test_finalize_presidio_anonymize_numbered_tokens_identical_spans():
     assert set(re.findall(r"<[A-Z_]+_\d+>", result)) == set(pii_tokens.keys())
 
 
+def test_finalize_presidio_anonymize_numbered_tokens_nested_higher_score_span():
+    import re
+
+    from litellm.proxy.guardrails.guardrail_hooks.presidio import _OPTIONAL_PresidioPIIMasking
+
+    guardrail = _OPTIONAL_PresidioPIIMasking(mock_testing=True)
+    text = "Card: 1234-5678-9012-3456."
+    analyze_results = [
+        {"entity_type": "CREDIT_CARD", "start": 6, "end": 25, "score": 0.70},
+        {"entity_type": "US_BANK_NUMBER", "start": 6, "end": 10, "score": 0.95},
+    ]
+    request_data = {"metadata": {}}
+    masked_entity_count = {}
+    result = guardrail._finalize_presidio_anonymize_numbered_tokens(
+        text=text,
+        analyze_results=analyze_results,
+        request_data=request_data,
+        masked_entity_count=masked_entity_count,
+    )
+    assert result == "Card: <CREDIT_CARD_1>."
+    pii_tokens = request_data["metadata"]["pii_tokens"]
+    assert pii_tokens["<CREDIT_CARD_1>"] == "1234-5678-9012-3456"
+    assert "<US_BANK_NUMBER_" not in pii_tokens
+    assert "-5678" not in result
+    assert set(re.findall(r"<[A-Z_]+_\d+>", result)) == set(pii_tokens.keys())
+
+
 def test_resolve_overlapping_spans_empty_and_disjoint():
     from litellm.proxy.guardrails.guardrail_hooks.presidio import _OPTIONAL_PresidioPIIMasking
 
-    assert _OPTIONAL_PresidioPIIMasking._resolve_overlapping_spans([]) == []
+    assert len(_OPTIONAL_PresidioPIIMasking._resolve_overlapping_spans([])) == 0
     disjoint = [
         {"entity_type": "PERSON", "start": 0, "end": 5, "score": 0.9},
         {"entity_type": "PHONE_NUMBER", "start": 10, "end": 20, "score": 0.8},
     ]
-    assert _OPTIONAL_PresidioPIIMasking._resolve_overlapping_spans(disjoint) == disjoint
+    resolved = _OPTIONAL_PresidioPIIMasking._resolve_overlapping_spans(disjoint)
+    assert len(resolved) == 2
+    assert resolved[0]["start"] == 0 and resolved[0]["end"] == 5
+    assert resolved[1]["start"] == 10 and resolved[1]["end"] == 20
 
 
 @pytest.mark.asyncio