diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index 1b53c5bad59..8d78393d687 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -68,6 +68,18 @@ class _PresidioAnonymizeResponse(TypedDict): items: ReadOnly[NotRequired[list[_PresidioAnonymizeItem]]] +_LoopSemaphores = dict[asyncio.AbstractEventLoop, asyncio.Semaphore] + + +def _json_escaped_len(text: str) -> int: + """ + Byte length of ``text`` as it appears serialized inside the JSON request + body sent to Presidio (``json.dumps`` escapes non-ASCII characters, so a + 3-byte UTF-8 character can occupy 6+ bytes on the wire). + """ + return len(json.dumps(text).encode("utf-8")) - 2 # strip the surrounding quotes + + class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): user_api_key_cache = None ad_hoc_recognizers: list[str] | None = None @@ -141,6 +153,10 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): # Loop-bound session cache for background threads self._loop_sessions: dict[asyncio.AbstractEventLoop, aiohttp.ClientSession] = {} + # Per-loop semaphores bounding chunked-analyze fan-out across ALL + # concurrent oversized blocks/requests on this instance, not per call + self._loop_chunk_semaphores: _LoopSemaphores = {} # mutable-ok: per-loop semaphore cache + if mock_testing is True: # for testing purposes only return @@ -302,7 +318,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): text and len(text) > 1 and self.mock_redacted_text is None - and len(text.encode("utf-8")) > self.presidio_analyze_chunk_size_bytes + and _json_escaped_len(text) > self.presidio_analyze_chunk_size_bytes ): return await self._analyze_text_chunked( text=text, @@ -434,9 +450,9 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): """ Analyze an oversized text by splitting it into overlapping chunks. - Each chunk is at most ``presidio_analyze_chunk_size_bytes`` UTF-8 bytes, - so every /analyze call stays below the analyzer deployment's request - body limit; per-chunk results are remapped onto the original text and + Each chunk serializes to at most ``presidio_analyze_chunk_size_bytes`` + bytes inside the JSON request body, so every /analyze call stays below + the analyzer deployment's request body limit; per-chunk results are remapped onto the original text and merged. Raises exactly like a single ``analyze_text`` call if any chunk fails. @@ -454,9 +470,14 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): self.presidio_analyze_chunk_size_bytes, len(text_chunks), ) - # Bound the fan-out so a single oversized request cannot saturate the - # analyzer; excess chunks wait here instead of piling onto the pool. - analyze_semaphore: Final = asyncio.Semaphore(PRESIDIO_ANALYZE_CHUNK_CONCURRENCY) + # Bound the fan-out so oversized requests cannot saturate the analyzer. + # The semaphore is shared per event loop across every chunked call on + # this instance, so many oversized blocks in one request (or many + # concurrent requests) still hold at most this many analyzer calls in + # flight. On the proxy's main thread the shared-session lock in + # _get_session_iterator additionally serializes the HTTP calls; the + # bound matters for loop-bound sessions (background threads). + analyze_semaphore: Final = self._get_chunk_semaphore() async def _analyze_chunk_bounded( chunk_text: str, @@ -479,9 +500,27 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): # analyze_text only returns a non-list shape when mock_redacted_text # is set, and the chunked path is never entered in that case. typed_result = cast("list[PresidioAnalyzeResponseItem]", result) # cast-ok: gather() erases element type - chunk_results.append(typed_result) + # Apply the configured score thresholds and deny list BEFORE the + # overlap merge: a below-threshold detection must not win overlap + # resolution against one the thresholds would keep. The same filter + # runs again downstream in check_pii, where it is a no-op for the + # already-filtered items. + filtered_result = self.filter_analyze_results_by_score(analyze_results=typed_result) + chunk_results.append( + cast("list[PresidioAnalyzeResponseItem]", filtered_result) # cast-ok: list input yields list + ) return self._merge_chunked_analyze_results(text_chunks=text_chunks, chunk_results=chunk_results) + def _get_chunk_semaphore(self) -> asyncio.Semaphore: + """Per-event-loop semaphore shared by all chunked analyze calls on this instance.""" + loop: Final = asyncio.get_running_loop() + existing: Final = self._loop_chunk_semaphores.get(loop) + if existing is not None: + return existing + created: Final = asyncio.Semaphore(PRESIDIO_ANALYZE_CHUNK_CONCURRENCY) + self._loop_chunk_semaphores[loop] = created + return created + @staticmethod def _coerce_analyze_chunk_size(value: int | None) -> int: """ @@ -489,9 +528,9 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): Non-positive values would either bypass chunking entirely or degenerate it into per-character splits (silently disabling detection), so they are - replaced by the default; values below 4 bytes (the widest UTF-8 - character) are floored to 4 so a single character always fits in a - chunk and the chunked path can never re-enter itself. + replaced by the default; values below 4 bytes are floored to 4 and the + splitter always emits at least one character per chunk, so the chunked + path can never re-enter itself. """ if not value or value <= 0: return DEFAULT_PRESIDIO_ANALYZE_CHUNK_SIZE_BYTES @@ -504,7 +543,10 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): overlap_chars: int, ) -> Sequence[tuple[int, str]]: """ - Split ``text`` into chunks of at most ``chunk_size_bytes`` UTF-8 bytes. + Split ``text`` into chunks whose JSON-serialized form is at most + ``chunk_size_bytes`` bytes (the analyzer body limit applies to the + JSON request body, where non-ASCII characters are escaped and larger + than their raw UTF-8 encoding). Consecutive chunks overlap by up to ``overlap_chars`` characters so a PII entity up to that length lying across a chunk boundary is still @@ -518,15 +560,23 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): text_len: Final = len(text) start = 0 # rebind-ok: chunk cursor advances across the loop while start < text_len: - # Byte-truncate a char-count-bounded slice, then drop the at most - # one trailing character the truncation split, so every chunk ends - # on a character boundary and holds at most chunk_size_bytes. + # Serialized length of a character is at least 1 byte, so a slice + # of chunk_size_bytes characters is a sufficient search window. candidate = text[start : start + chunk_size_bytes] - chunk = candidate.encode("utf-8")[:chunk_size_bytes].decode("utf-8", errors="ignore") - if not chunk: - # chunk_size_bytes is below one character's UTF-8 width; emit a - # single character rather than an empty chunk. - chunk = candidate[:1] + if _json_escaped_len(candidate) <= chunk_size_bytes: + chunk = candidate + else: + # Largest prefix whose serialized form fits the budget. + low, high = 1, len(candidate) + while low < high: + mid = (low + high + 1) // 2 + if _json_escaped_len(candidate[:mid]) <= chunk_size_bytes: + low = mid + else: + high = mid - 1 + # low >= 1 keeps the loop advancing even when a single + # character serializes over a (floored, tiny) budget. + chunk = candidate[:low] end = start + len(chunk) chunks.append((start, chunk)) if end >= text_len: diff --git a/litellm/proxy/guardrails/guardrail_initializers.py b/litellm/proxy/guardrails/guardrail_initializers.py index c07daeea082..595fade8387 100644 --- a/litellm/proxy/guardrails/guardrail_initializers.py +++ b/litellm/proxy/guardrails/guardrail_initializers.py @@ -100,11 +100,15 @@ def initialize_presidio(litellm_params: LitellmParams, guardrail: Guardrail): presidio_anonymizer_api_base=litellm_params.presidio_anonymizer_api_base, presidio_language=litellm_params.presidio_language, presidio_entities_deny_list=litellm_params.presidio_entities_deny_list, - presidio_analyze_chunk_size_bytes=litellm_params.presidio_analyze_chunk_size_bytes, apply_to_output=False, ) params.update(overrides) - callback: Final = _OPTIONAL_PresidioPIIMasking(**params) + # Passed outside the heterogeneous params dict so the argument keeps + # its precise int | None type. + callback: Final = _OPTIONAL_PresidioPIIMasking( + presidio_analyze_chunk_size_bytes=litellm_params.presidio_analyze_chunk_size_bytes, + **params, + ) litellm.logging_callback_manager.add_litellm_callback(callback) return callback 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 647704b0c42..4ee6741ee02 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py @@ -22,9 +22,7 @@ from litellm.types.utils import Choices, Message, ModelResponse from litellm.exceptions import BlockedPiiEntityError -def _make_mock_session_iterator( - json_response, status=200, content_type="application/json", text_response="" -): +def _make_mock_session_iterator(json_response, status=200, content_type="application/json", text_response=""): """Create a mock _get_session_iterator that yields a session returning json_response.""" @asynccontextmanager @@ -100,9 +98,7 @@ def mock_cache(): @pytest.mark.asyncio -async def test_multimodal_message_format_completion_call_type( - presidio_guardrail, mock_user_api_key, mock_cache -): +async def test_multimodal_message_format_completion_call_type(presidio_guardrail, mock_user_api_key, mock_cache): """ Test Presidio PII masking with multimodal message format (content as list) for completion call type. @@ -247,9 +243,7 @@ async def test_multimodal_message_format_anthropic_messages_call_type( @pytest.mark.asyncio -async def test_multimodal_message_multiple_content_items( - presidio_guardrail, mock_user_api_key, mock_cache -): +async def test_multimodal_message_multiple_content_items(presidio_guardrail, mock_user_api_key, mock_cache): """ Test Presidio PII masking with multiple content items in the content list. """ @@ -303,9 +297,7 @@ async def test_multimodal_message_multiple_content_items( @pytest.mark.asyncio -async def test_mixed_string_and_list_content( - presidio_guardrail, mock_user_api_key, mock_cache -): +async def test_mixed_string_and_list_content(presidio_guardrail, mock_user_api_key, mock_cache): """ Test Presidio PII masking with mixed string and list content formats. """ @@ -370,9 +362,7 @@ async def test_mixed_string_and_list_content( @pytest.mark.asyncio -async def test_content_list_without_text_field( - presidio_guardrail, mock_user_api_key, mock_cache -): +async def test_content_list_without_text_field(presidio_guardrail, mock_user_api_key, mock_cache): """ Test Presidio PII masking gracefully handles content items without text field (e.g., image content items). @@ -629,9 +619,7 @@ async def test_logging_hook_masks_the_response_too(presidio_guardrail): @pytest.mark.asyncio -async def test_logging_only_does_not_mask_pre_call_request( - mock_user_api_key, mock_cache -): +async def test_logging_only_does_not_mask_pre_call_request(mock_user_api_key, mock_cache): """ A guardrail configured with `logging_only` must only mask PII for logs/traces, never for the request sent to the model. `async_pre_call_hook` should leave the @@ -718,9 +706,7 @@ async def test_presidio_sets_guardrail_information_in_request_data(): assert "metadata" in request_data assert "standard_logging_guardrail_information" in request_data["metadata"] - guardrail_info_list = request_data["metadata"][ - "standard_logging_guardrail_information" - ] + guardrail_info_list = request_data["metadata"]["standard_logging_guardrail_information"] assert isinstance(guardrail_info_list, list) assert len(guardrail_info_list) > 0 @@ -847,20 +833,14 @@ async def test_presidio_filter_scope_initializer(monkeypatch): import litellm.proxy.guardrails.guardrail_hooks.presidio as presidio_mod import litellm.proxy.guardrails.guardrail_initializers as gi - monkeypatch.setattr( - presidio_mod, "_OPTIONAL_PresidioPIIMasking", DummyGuardrail, raising=False - ) - monkeypatch.setattr( - gi, "_OPTIONAL_PresidioPIIMasking", DummyGuardrail, raising=False - ) + monkeypatch.setattr(presidio_mod, "_OPTIONAL_PresidioPIIMasking", DummyGuardrail, raising=False) + monkeypatch.setattr(gi, "_OPTIONAL_PresidioPIIMasking", DummyGuardrail, raising=False) # input-only created.clear() from litellm.proxy.guardrails.guardrail_initializers import initialize_presidio - params_input = LitellmParams( - guardrail="presidio", mode="pre_call", presidio_filter_scope="input" - ) + params_input = LitellmParams(guardrail="presidio", mode="pre_call", presidio_filter_scope="input") guardrail_dict = {"guardrail_name": "g1"} cb = initialize_presidio(params_input, guardrail_dict) assert cb is created[0] @@ -868,18 +848,14 @@ async def test_presidio_filter_scope_initializer(monkeypatch): # output-only created.clear() - params_output = LitellmParams( - guardrail="presidio", mode="pre_call", presidio_filter_scope="output" - ) + params_output = LitellmParams(guardrail="presidio", mode="pre_call", presidio_filter_scope="output") cb = initialize_presidio(params_output, guardrail_dict) assert len(created) == 1 assert created[0].apply_to_output is True # both -> expect two callbacks (input + output) created.clear() - params_both = LitellmParams( - guardrail="presidio", mode="pre_call", presidio_filter_scope="both" - ) + params_both = LitellmParams(guardrail="presidio", mode="pre_call", presidio_filter_scope="both") cb = initialize_presidio(params_both, guardrail_dict) assert len(created) == 2 assert any(not c.apply_to_output for c in created) @@ -887,9 +863,7 @@ async def test_presidio_filter_scope_initializer(monkeypatch): @pytest.mark.asyncio -async def test_empty_content_handling( - presidio_guardrail, mock_user_api_key, mock_cache -): +async def test_empty_content_handling(presidio_guardrail, mock_user_api_key, mock_cache): """ Test that Presidio handles empty content gracefully. @@ -945,9 +919,7 @@ async def test_empty_content_handling( @pytest.mark.asyncio -async def test_whitespace_only_content( - presidio_guardrail, mock_user_api_key, mock_cache -): +async def test_whitespace_only_content(presidio_guardrail, mock_user_api_key, mock_cache): """ Test that Presidio handles whitespace-only content gracefully. @@ -1142,9 +1114,7 @@ async def test_analyze_text_list_with_non_dict_items(): "invalid_string_item", {"entity_type": "EMAIL", "start": 10, "end": 25, "score": 0.85}, ] - with patch.object( - presidio, "_get_session_iterator", _make_mock_session_iterator(json_response) - ): + with patch.object(presidio, "_get_session_iterator", _make_mock_session_iterator(json_response)): result = await presidio.analyze_text( text="some text", presidio_config=None, @@ -1156,9 +1126,7 @@ async def test_analyze_text_list_with_non_dict_items(): @pytest.mark.asyncio -async def test_tool_calling_complete_scenario( - presidio_guardrail, mock_user_api_key, mock_cache -): +async def test_tool_calling_complete_scenario(presidio_guardrail, mock_user_api_key, mock_cache): """ Test complete tool calling scenario with PII in user message. @@ -1224,9 +1192,7 @@ def test_filter_drops_low_score_detection(): mock_testing=True, presidio_score_thresholds={PiiEntityType.CREDIT_CARD: 0.8}, ) - analyze_results = [ - {"entity_type": PiiEntityType.CREDIT_CARD, "score": 0.7, "start": 0, "end": 4} - ] + analyze_results = [{"entity_type": PiiEntityType.CREDIT_CARD, "score": 0.7, "start": 0, "end": 4}] filtered = guardrail.filter_analyze_results_by_score(analyze_results) assert filtered == [] @@ -1240,9 +1206,7 @@ def test_filter_preserves_high_score_detection(): mock_testing=True, presidio_score_thresholds={PiiEntityType.CREDIT_CARD: 0.8}, ) - analyze_results = [ - {"entity_type": PiiEntityType.CREDIT_CARD, "score": 0.9, "start": 0, "end": 4} - ] + analyze_results = [{"entity_type": PiiEntityType.CREDIT_CARD, "score": 0.9, "start": 0, "end": 4}] filtered = guardrail.filter_analyze_results_by_score(analyze_results) assert len(filtered) == 1 @@ -1379,15 +1343,11 @@ def test_blocking_respects_threshold_filter(): presidio_score_thresholds={PiiEntityType.CREDIT_CARD: 0.9}, ) - low_score_results = [ - {"entity_type": PiiEntityType.CREDIT_CARD, "score": 0.7, "start": 0, "end": 4} - ] + low_score_results = [{"entity_type": PiiEntityType.CREDIT_CARD, "score": 0.7, "start": 0, "end": 4}] filtered = guardrail.filter_analyze_results_by_score(low_score_results) guardrail.raise_exception_if_blocked_entities_detected(filtered) - high_score_results = [ - {"entity_type": PiiEntityType.CREDIT_CARD, "score": 0.95, "start": 0, "end": 4} - ] + high_score_results = [{"entity_type": PiiEntityType.CREDIT_CARD, "score": 0.95, "start": 0, "end": 4}] filtered_high = guardrail.filter_analyze_results_by_score(high_score_results) with pytest.raises(BlockedPiiEntityError): guardrail.raise_exception_if_blocked_entities_detected(filtered_high) @@ -1448,9 +1408,7 @@ async def test_get_session_iterator_thread_safety(presidio_guardrail): # Run the background thread test bg_future = asyncio.Future() - t = threading.Thread( - target=thread_target, args=(asyncio.get_running_loop(), bg_future) - ) + t = threading.Thread(target=thread_target, args=(asyncio.get_running_loop(), bg_future)) t.start() t.join() @@ -1659,9 +1617,7 @@ async def test_anonymize_text_non_json_content_type(): ) with patch.object(guardrail, "_get_session_iterator", mock_iterator): - with pytest.raises( - Exception, match="Presidio anonymizer returned non-JSON Content-Type" - ): + with pytest.raises(Exception, match="Presidio anonymizer returned non-JSON Content-Type"): await guardrail.anonymize_text( text="Hello world", analyze_results=[{"start": 0, "end": 5, "entity_type": "PERSON"}], @@ -1719,9 +1675,7 @@ async def test_pii_tokens_stored_in_metadata_not_top_level(presidio_guardrail): mock_cache = DualCache() test_data = { - "messages": [ - {"role": "user", "content": "My name is John and my phone is 555-123-4567"} - ], + "messages": [{"role": "user", "content": "My name is John and my phone is 555-123-4567"}], "model": "claude-haiku-4-5-20251001", "metadata": {}, } @@ -1870,9 +1824,7 @@ async def test_metadata_none_does_not_crash(): ) # No pii_tokens to unmask, so content stays as-is - assert ( - response.choices[0].message.content == f"Hello {token_key}, how can I help you?" - ) + assert response.choices[0].message.content == f"Hello {token_key}, how can I help you?" # --------------------------------------------------------------------------- @@ -2049,9 +2001,7 @@ async def test_anthropic_native_response_unmasking(): response=anthropic_response, ) - assert result["content"][0]["text"] == ( - "Hello John Smith, your number is 555-123-4567." - ) + assert result["content"][0]["text"] == ("Hello John Smith, your number is 555-123-4567.") @pytest.mark.asyncio @@ -2170,9 +2120,7 @@ async def test_streaming_bytes_chunks_are_yielded_not_discarded(): ): chunks.append(chunk) - assert any( - isinstance(c, bytes) for c in chunks - ), "bytes chunks must not be discarded" + assert any(isinstance(c, bytes) for c in chunks), "bytes chunks must not be discarded" assert byte_chunk in chunks @@ -2282,9 +2230,7 @@ async def test_apply_to_output_streaming_mixed_chunks_flushes_and_warns(): mock_user_api_key = UserAPIKeyAuth(api_key="test-key") received = [] - with patch( - "litellm.proxy.guardrails.guardrail_hooks.presidio.verbose_proxy_logger" - ) as mock_logger: + with patch("litellm.proxy.guardrails.guardrail_hooks.presidio.verbose_proxy_logger") as mock_logger: async for chunk in guardrail.async_post_call_streaming_iterator_hook( user_api_key_dict=mock_user_api_key, response=mock_stream(), @@ -2396,9 +2342,7 @@ async def test_apply_to_output_streaming_bytes_only_logs_warning(): mock_user_api_key = UserAPIKeyAuth(api_key="test-key") collected = [] - with patch( - "litellm.proxy.guardrails.guardrail_hooks.presidio.verbose_proxy_logger" - ) as mock_logger: + with patch("litellm.proxy.guardrails.guardrail_hooks.presidio.verbose_proxy_logger") as mock_logger: async for chunk in guardrail.async_post_call_streaming_iterator_hook( user_api_key_dict=mock_user_api_key, response=mock_stream(), @@ -2521,10 +2465,7 @@ async def test_output_parse_pii_streaming_responses_completed_event_unmasked( collected.append(chunk) assert collected == [completed_event] - assert ( - collected[0].response.output[0].content[0].text - == "Reach me at john@example.com today." - ) + assert collected[0].response.output[0].content[0].text == "Reach me at john@example.com today." @pytest.mark.asyncio @@ -2587,9 +2528,7 @@ async def test_anonymize_text_uses_correct_positions_no_parse_pii(): original text using those positions, which produces garbled output with remnants of original PII data. """ - original_text = ( - "My name is John Smith, my email is john@example.com, phone 555-867-5309" - ) + original_text = "My name is John Smith, my email is john@example.com, phone 555-867-5309" # Positions as returned by the analyzer (reference original text) analyze_results = [ {"end": 51, "entity_type": "EMAIL_ADDRESS", "score": 1.0, "start": 35}, @@ -2644,9 +2583,9 @@ async def test_anonymize_text_uses_correct_positions_no_parse_pii(): ) expected = "My name is , my email is , phone " - assert ( - result == expected - ), f"anonymize_text produced garbled output with PII remnants.\nExpected: {expected!r}\nGot: {result!r}" + assert result == expected, ( + f"anonymize_text produced garbled output with PII remnants.\nExpected: {expected!r}\nGot: {result!r}" + ) assert masked_entity_count == { "PERSON": 1, "EMAIL_ADDRESS": 1, @@ -2665,9 +2604,7 @@ async def test_anonymize_text_uses_correct_positions_with_parse_pii(): tokens and the pii_tokens mapping, not positions from anonymizer items (which reference the anonymized output text). """ - original_text = ( - "My name is John Smith, my email is john@example.com, phone 555-867-5309" - ) + original_text = "My name is John Smith, my email is john@example.com, phone 555-867-5309" analyze_results = [ {"end": 51, "entity_type": "EMAIL_ADDRESS", "score": 1.0, "start": 35}, {"end": 21, "entity_type": "PERSON", "score": 0.85, "start": 11}, @@ -2783,17 +2720,13 @@ def test_unmask_sse_bytes_chunk_ignores_non_text_delta(): def test_unmask_sse_bytes_chunk_handles_malformed_json(): chunk = b"data: {not valid json}\n\n" - result = _OPTIONAL_PresidioPIIMasking._unmask_sse_bytes_chunk( - chunk, {"": "Bobby"} - ) + result = _OPTIONAL_PresidioPIIMasking._unmask_sse_bytes_chunk(chunk, {"": "Bobby"}) assert result == chunk def test_unmask_sse_bytes_chunk_handles_unicode_decode_error(): chunk = b"\xff\xfe invalid utf-8" - result = _OPTIONAL_PresidioPIIMasking._unmask_sse_bytes_chunk( - chunk, {"": "Bobby"} - ) + result = _OPTIONAL_PresidioPIIMasking._unmask_sse_bytes_chunk(chunk, {"": "Bobby"}) assert result == chunk @@ -2827,9 +2760,7 @@ def test_unmask_sse_bytes_chunk_handles_crlf_line_endings(): } crlf_chunk = ("data: " + json.dumps(event) + "\r\ndata: [DONE]\r\n").encode("utf-8") - result = _OPTIONAL_PresidioPIIMasking._unmask_sse_bytes_chunk( - crlf_chunk, pii_tokens - ) + result = _OPTIONAL_PresidioPIIMasking._unmask_sse_bytes_chunk(crlf_chunk, pii_tokens) decoded = result.decode("utf-8") parsed = json.loads(decoded.split("data: ", 1)[1].split("\n")[0].strip()) @@ -2948,10 +2879,7 @@ def _make_marker_session_iterator( if url.endswith("analyze"): recorded_analyze_payloads.append(payload) text = payload["text"] - if ( - analyzer_body_limit_bytes is not None - and len(text.encode("utf-8")) > analyzer_body_limit_bytes - ): + if analyzer_body_limit_bytes is not None and len(text.encode("utf-8")) > analyzer_body_limit_bytes: return MockResponse( 413, { @@ -2973,9 +2901,7 @@ def _make_marker_session_iterator( if recorded_anonymize_payloads is not None: recorded_anonymize_payloads.append(payload) text = payload["text"] - items = sorted( - payload["analyzer_results"], key=lambda r: r["start"], reverse=True - ) + items = sorted(payload["analyzer_results"], key=lambda r: r["start"], reverse=True) for r in items: text = text[: r["start"]] + "<" + r["entity_type"] + ">" + text[r["end"] :] return MockResponse( @@ -3015,9 +2941,7 @@ def _oversized_marker_text(): def test_split_text_for_analysis_offsets_and_byte_budget(): text = " ".join(f"word{i}" for i in range(200)) - chunks = _OPTIONAL_PresidioPIIMasking._split_text_for_analysis( - text=text, chunk_size_bytes=100, overlap_chars=20 - ) + chunks = _OPTIONAL_PresidioPIIMasking._split_text_for_analysis(text=text, chunk_size_bytes=100, overlap_chars=20) assert len(chunks) > 1 for offset, chunk in chunks: assert len(chunk.encode("utf-8")) <= 100 @@ -3032,9 +2956,7 @@ def test_split_text_for_analysis_offsets_and_byte_budget(): def test_split_text_for_analysis_multibyte_characters(): text = "émoji🙂 çafé " * 120 - chunks = _OPTIONAL_PresidioPIIMasking._split_text_for_analysis( - text=text, chunk_size_bytes=64, overlap_chars=8 - ) + chunks = _OPTIONAL_PresidioPIIMasking._split_text_for_analysis(text=text, chunk_size_bytes=64, overlap_chars=8) assert len(chunks) > 1 for offset, chunk in chunks: assert len(chunk.encode("utf-8")) <= 64 @@ -3044,9 +2966,7 @@ def test_split_text_for_analysis_multibyte_characters(): def test_split_text_for_analysis_under_budget_returns_single_chunk(): text = "short text" - chunks = _OPTIONAL_PresidioPIIMasking._split_text_for_analysis( - text=text, chunk_size_bytes=100, overlap_chars=20 - ) + chunks = _OPTIONAL_PresidioPIIMasking._split_text_for_analysis(text=text, chunk_size_bytes=100, overlap_chars=20) assert chunks == [(0, text)] @@ -3055,12 +2975,8 @@ async def test_analyze_text_single_call_when_under_limit(): guardrail = _chunking_guardrail(chunk_size_bytes=10_000) payloads = [] text = f"my card is {CHUNK_MARKER_ONE} thanks" - with patch.object( - guardrail, "_get_session_iterator", _make_marker_session_iterator(payloads) - ): - results = await guardrail.analyze_text( - text=text, presidio_config=None, request_data={} - ) + with patch.object(guardrail, "_get_session_iterator", _make_marker_session_iterator(payloads)): + results = await guardrail.analyze_text(text=text, presidio_config=None, request_data={}) assert len(payloads) == 1 assert payloads[0]["text"] == text assert len(results) == 1 @@ -3088,9 +3004,7 @@ async def test_analyze_text_chunks_oversized_text_and_remaps_offsets(): "_get_session_iterator", _make_marker_session_iterator(payloads, analyzer_body_limit_bytes=100), ): - results = await guardrail.analyze_text( - text=text, presidio_config=None, request_data={} - ) + results = await guardrail.analyze_text(text=text, presidio_config=None, request_data={}) assert len(payloads) > 1 for payload in payloads: assert len(payload["text"].encode("utf-8")) <= 100 @@ -3118,9 +3032,7 @@ async def test_check_pii_masks_oversized_text_with_chunking(): recorded_anonymize_payloads=anonymize_payloads, ), ): - masked = await guardrail.check_pii( - text=text, output_parse_pii=False, presidio_config=None, request_data={} - ) + masked = await guardrail.check_pii(text=text, output_parse_pii=False, presidio_config=None, request_data={}) assert CHUNK_MARKER_ONE not in masked assert CHUNK_MARKER_TWO not in masked assert masked.count("") == 2 @@ -3177,9 +3089,7 @@ async def test_analyze_text_chunked_failure_stays_fail_closed(): _make_marker_session_iterator(payloads, analyzer_body_limit_bytes=10), ): with pytest.raises(GuardrailRaisedException, match="HTTP 413"): - await guardrail.analyze_text( - text=text, presidio_config=None, request_data={} - ) + await guardrail.analyze_text(text=text, presidio_config=None, request_data={}) def test_presidio_analyze_chunk_size_default_and_validation(): @@ -3188,14 +3098,10 @@ def test_presidio_analyze_chunk_size_default_and_validation(): guardrail = _OPTIONAL_PresidioPIIMasking(mock_testing=True) assert guardrail.presidio_analyze_chunk_size_bytes == DEFAULT_PRESIDIO_ANALYZE_CHUNK_SIZE_BYTES - nonpositive = _OPTIONAL_PresidioPIIMasking( - mock_testing=True, presidio_analyze_chunk_size_bytes=-5 - ) + nonpositive = _OPTIONAL_PresidioPIIMasking(mock_testing=True, presidio_analyze_chunk_size_bytes=-5) assert nonpositive.presidio_analyze_chunk_size_bytes == DEFAULT_PRESIDIO_ANALYZE_CHUNK_SIZE_BYTES - custom = _OPTIONAL_PresidioPIIMasking( - mock_testing=True, presidio_analyze_chunk_size_bytes=1234 - ) + custom = _OPTIONAL_PresidioPIIMasking(mock_testing=True, presidio_analyze_chunk_size_bytes=1234) assert custom.presidio_analyze_chunk_size_bytes == 1234 @@ -3276,9 +3182,7 @@ async def test_tiny_chunk_size_with_multibyte_text_terminates(): guardrail = _chunking_guardrail(chunk_size_bytes=1) assert guardrail.presidio_analyze_chunk_size_bytes == 4 payloads = [] - with patch.object( - guardrail, "_get_session_iterator", _make_marker_session_iterator(payloads) - ): + with patch.object(guardrail, "_get_session_iterator", _make_marker_session_iterator(payloads)): results = await guardrail.analyze_text( text="\U0001f642\U0001f642\U0001f642ab", presidio_config=None, request_data={} ) @@ -3332,3 +3236,147 @@ async def test_chunked_analyze_concurrency_is_bounded(): await guardrail.analyze_text(text="a" * 400, presidio_config=None, request_data={}) assert state["peak"] >= 2 assert state["peak"] <= PRESIDIO_ANALYZE_CHUNK_CONCURRENCY + + +def test_split_text_accounts_for_json_body_expansion(): + """Non-ASCII text expands under JSON escaping; the budget must apply to the + serialized form or a chunk can still exceed the analyzer body limit.""" + import json as json_module + + text = "これは個人情報テストです。" * 200 # 3-byte UTF-8 chars, 6-byte escapes + budget = 1000 + chunks = _OPTIONAL_PresidioPIIMasking._split_text_for_analysis(text=text, chunk_size_bytes=budget, overlap_chars=8) + assert len(chunks) > 1 + for offset, chunk in chunks: + assert len(json_module.dumps(chunk).encode("utf-8")) - 2 <= budget + assert text[offset : offset + len(chunk)] == chunk + # full coverage: last chunk reaches the end of the text + last_offset, last_chunk = chunks[-1] + assert last_offset + len(last_chunk) == len(text) + + +@pytest.mark.asyncio +async def test_chunked_analyze_applies_score_threshold_before_merge(): + """A below-threshold long span must not win overlap resolution against an + above-threshold detection of the same type (it would then be dropped by the + downstream threshold filter, leaving the entity unmasked).""" + guardrail = _chunking_guardrail( + chunk_size_bytes=100, + presidio_score_thresholds={"CREDIT_CARD": 0.6}, + ) + marker_text = "x" * 40 + CHUNK_MARKER_ONE + "x" * 80 # single chunked text + + @asynccontextmanager + async def mock_iterator(): + class MockResponse: + status = 200 + content_type = "application/json" + headers = {"Content-Type": "application/json"} + + def __init__(self, body): + self._body = body + + async def text(self): + import json as json_module + + return json_module.dumps(self._body) + + async def json(self): + return self._body + + async def __aenter__(self): + return self + + async def __aexit__(self, *args): + return False + + class MockSession: + def post(self, url, json=None, headers=None): + text = json["text"] + idx = text.find(CHUNK_MARKER_ONE) + if idx == -1: + return MockResponse([]) + return MockResponse( + [ + # long, below-threshold span engulfing the marker + { + "entity_type": "CREDIT_CARD", + "start": max(idx - 5, 0), + "end": idx + len(CHUNK_MARKER_ONE) + 5, + "score": 0.3, + }, + # the true, above-threshold detection + { + "entity_type": "CREDIT_CARD", + "start": idx, + "end": idx + len(CHUNK_MARKER_ONE), + "score": 0.9, + }, + ] + ) + + async def __aenter__(self): + return self + + async def __aexit__(self, *args): + return False + + yield MockSession() + + with patch.object(guardrail, "_get_session_iterator", mock_iterator): + results = await guardrail.analyze_text(text=marker_text, presidio_config=None, request_data={}) + kept = [r for r in results if r.get("entity_type") == "CREDIT_CARD"] + assert any(r.get("score") == 0.9 for r in kept), kept + assert all(r.get("score") != 0.3 for r in kept), kept + + +@pytest.mark.asyncio +async def test_chunk_fanout_bound_is_shared_across_concurrent_calls(): + """The chunk semaphore is per event loop and instance, so several oversized + blocks analyzed concurrently share ONE bound instead of getting 8 each.""" + from litellm.constants import PRESIDIO_ANALYZE_CHUNK_CONCURRENCY + + guardrail = _chunking_guardrail(chunk_size_bytes=10) + state = {"active": 0, "peak": 0} + + @asynccontextmanager + async def mock_iterator(): + class MockResponse: + status = 200 + content_type = "application/json" + headers = {"Content-Type": "application/json"} + + async def text(self): + return "[]" + + async def json(self): + state["active"] += 1 + state["peak"] = max(state["peak"], state["active"]) + await asyncio.sleep(0.005) + state["active"] -= 1 + return [] + + async def __aenter__(self): + return self + + async def __aexit__(self, *args): + return False + + class MockSession: + def post(self, url, json=None, headers=None): + return MockResponse() + + async def __aenter__(self): + return self + + async def __aexit__(self, *args): + return False + + yield MockSession() + + with patch.object(guardrail, "_get_session_iterator", mock_iterator): + await asyncio.gather( + *(guardrail.analyze_text(text="a" * 400, presidio_config=None, request_data={}) for _ in range(4)) + ) + assert state["peak"] >= 2 + assert state["peak"] <= PRESIDIO_ANALYZE_CHUNK_CONCURRENCY