From 3891cdc2a3c3c8b810cc80faedb0cec1f1a61050 Mon Sep 17 00:00:00 2001 From: musi22 Date: Sun, 20 Sep 2026 19:36:03 +0530 Subject: [PATCH] test(presidio): add regression tests for overlapping entity spans (#42130) --- .../guardrails/guardrail_hooks/presidio.py | 33 ++++++- .../guardrail_hooks/test_presidio.py | 93 +++++++++++++++++++ 2 files changed, 125 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index 70ea21320ee..57e55d16aef 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -693,6 +693,35 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): masked_entity_count[entity_type] = masked_entity_count.get(entity_type, 0) + 1 return redacted_text["text"] + @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), + ), + ) + 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) + + return sorted(accepted_spans, key=lambda x: x["start"]) + def _finalize_presidio_anonymize_numbered_tokens( self, text: str, @@ -718,9 +747,11 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): request_data["metadata"]["pii_tokens"] = {} pii_tokens: Final = request_data["metadata"]["pii_tokens"] + valid_analyze_results = self._resolve_overlapping_spans(analyze_results) + # Assign sequence numbers in forward (left-to-right) order so # that is the first entity in the text, etc. - sorted_forward: Final = sorted(analyze_results, key=lambda x: x["start"]) + sorted_forward: Final = sorted(valid_analyze_results, key=lambda x: x["start"]) seq_map: Final = {} for idx, ar in enumerate(sorted_forward, start=1): seq_map[(ar["start"], ar["end"])] = idx 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 33614d2eeca..7ee51c052b3 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py @@ -3114,6 +3114,99 @@ async def test_output_parse_pii_numbered_tokens_across_chunks(): assert pii_tokens[""] == CHUNK_MARKER_TWO +def test_finalize_presidio_anonymize_numbered_tokens_overlapping_spans(): + import re + + from litellm.proxy.guardrails.guardrail_hooks.presidio import _OPTIONAL_PresidioPIIMasking + + guardrail = _OPTIONAL_PresidioPIIMasking(mock_testing=True) + text = "My name is John Smith, phone +1 415 555 2671, please check the order." + analyze_results = [ + {"entity_type": "PERSON", "start": 11, "end": 21, "score": 0.9}, + {"entity_type": "PHONE_NUMBER", "start": 29, "end": 44, "score": 0.85}, + {"entity_type": "UK_NHS", "start": 32, "end": 44, "score": 0.6}, + ] + 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 == "My name is , phone , please check the order." + pii_tokens = request_data["metadata"]["pii_tokens"] + assert pii_tokens[""] == "John Smith" + assert pii_tokens[""] == "+1 415 555 2671" + assert "", result)) == set(pii_tokens.keys()) + + +def test_finalize_presidio_anonymize_numbered_tokens_three_way_overlap(): + import re + + from litellm.proxy.guardrails.guardrail_hooks.presidio import _OPTIONAL_PresidioPIIMasking + + guardrail = _OPTIONAL_PresidioPIIMasking(mock_testing=True) + text = "Contact: +1 415 555 2671, please call soon." + analyze_results = [ + {"entity_type": "PHONE_NUMBER", "start": 9, "end": 24, "score": 0.85}, + {"entity_type": "US_BANK_NUMBER", "start": 12, "end": 24, "score": 0.75}, + {"entity_type": "UK_NHS", "start": 16, "end": 24, "score": 0.6}, + ] + 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 == "Contact: , please call soon." + pii_tokens = request_data["metadata"]["pii_tokens"] + assert pii_tokens[""] == "+1 415 555 2671" + assert "", result)) == set(pii_tokens.keys()) + + +def test_finalize_presidio_anonymize_numbered_tokens_identical_spans(): + import re + + from litellm.proxy.guardrails.guardrail_hooks.presidio import _OPTIONAL_PresidioPIIMasking + + guardrail = _OPTIONAL_PresidioPIIMasking(mock_testing=True) + text = "Account: 9876543210 status." + analyze_results = [ + {"entity_type": "PHONE_NUMBER", "start": 9, "end": 19, "score": 0.85}, + {"entity_type": "US_BANK_NUMBER", "start": 9, "end": 19, "score": 0.70}, + ] + 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 == "Account: status." + pii_tokens = request_data["metadata"]["pii_tokens"] + assert pii_tokens[""] == "9876543210" + assert "", 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([]) == [] + 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 + + @pytest.mark.asyncio async def test_analyze_text_chunked_failure_stays_fail_closed(): """If one chunk still fails, the chunked path raises exactly like a single