diff --git a/litellm/constants.py b/litellm/constants.py index ed89474600e..084371e1dc0 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -296,6 +296,9 @@ GUARDRAIL_SCANNED_MESSAGES_CACHE_TTL_SECONDS: Final = int( os.getenv("GUARDRAIL_SCANNED_MESSAGES_CACHE_TTL_SECONDS", 24 * 60 * 60) ) BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS: Final = 25_000 +DEFAULT_PRESIDIO_ANALYZE_CHUNK_SIZE_BYTES: Final = 500_000 +PRESIDIO_ANALYZE_CHUNK_OVERLAP_CHARS: Final = 4096 +PRESIDIO_ANALYZE_CHUNK_CONCURRENCY: Final = 8 # Aggregation threshold: default to 80% of the asyncio queue maxsize so the check can always trigger. # Must be < LITELLM_ASYNCIO_QUEUE_MAXSIZE; if set higher the aggregation logic will never fire. MAX_SIZE_IN_MEMORY_QUEUE: Final = int(os.getenv("MAX_SIZE_IN_MEMORY_QUEUE", int(LITELLM_ASYNCIO_QUEUE_MAXSIZE * 0.8))) diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index bcee45355e3..8d78393d687 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -11,7 +11,7 @@ import asyncio import json import threading -from collections.abc import AsyncGenerator +from collections.abc import AsyncGenerator, Sequence from contextlib import asynccontextmanager from datetime import datetime from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypedDict, cast @@ -22,6 +22,11 @@ from typing_extensions import NotRequired, ReadOnly import litellm from litellm import get_secret from litellm._logging import verbose_proxy_logger +from litellm.constants import ( + DEFAULT_PRESIDIO_ANALYZE_CHUNK_SIZE_BYTES, + PRESIDIO_ANALYZE_CHUNK_CONCURRENCY, + PRESIDIO_ANALYZE_CHUNK_OVERLAP_CHARS, +) from litellm.types.utils import GenericGuardrailAPIInputs if TYPE_CHECKING: @@ -63,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 @@ -93,6 +110,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): presidio_language: str | None = None, presidio_score_thresholds: dict[PiiEntityType | str, float] | None = None, presidio_entities_deny_list: list[PiiEntityType | str] | None = None, + presidio_analyze_chunk_size_bytes: int | None = None, **kwargs, ): if logging_only is True: @@ -121,6 +139,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): self.presidio_score_thresholds: dict[PiiEntityType | str, float] = presidio_score_thresholds or {} self.presidio_entities_deny_list: list[PiiEntityType | str] = presidio_entities_deny_list or [] self.presidio_language = presidio_language or "en" + self.presidio_analyze_chunk_size_bytes: int = self._coerce_analyze_chunk_size(presidio_analyze_chunk_size_bytes) # Shared HTTP session to prevent memory leaks (issue #14540) self._http_session: aiohttp.ClientSession | None = None # Lock to prevent race conditions when creating session under concurrent load @@ -134,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 @@ -280,7 +303,28 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): ) -> list[PresidioAnalyzeResponseItem] | _PresidioAnonymizeResponse: """ Send text to the Presidio analyzer endpoint and get analysis results + + Texts larger than ``presidio_analyze_chunk_size_bytes`` (UTF-8) are split + into overlapping chunks, analyzed per chunk, and the per-chunk results + are remapped onto the original text. Presidio analyzer deployments + commonly cap the /analyze request body size (e.g. at 1 MB), and analyzer + latency grows with payload size. """ + # Chunk oversized texts before the try block so that a failing chunk + # keeps the same sanitized error message a single call would produce. + # A single-character text can never be split further, so it always + # takes the single-call path regardless of its encoded width. + if ( + text + and len(text) > 1 + and self.mock_redacted_text is None + and _json_escaped_len(text) > self.presidio_analyze_chunk_size_bytes + ): + return await self._analyze_text_chunked( + text=text, + presidio_config=presidio_config, + request_data=request_data, + ) try: # Skip empty or whitespace-only text to avoid Presidio errors # Common in tool/function calling where assistant content is empty @@ -397,6 +441,201 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): # contain API keys or other secrets) in error responses. raise Exception(f"Presidio PII analysis failed: {type(e).__name__}") from e + async def _analyze_text_chunked( + self, + text: str, + presidio_config: PresidioPerRequestConfig | None, + request_data: dict, # mutable-ok: shared per-request state dict, matching analyze_text's parameter + ) -> list[PresidioAnalyzeResponseItem]: # mutable-ok: analyze_text's declared return type requires list + """ + Analyze an oversized text by splitting it into overlapping chunks. + + 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. + + Only the analyzer side is chunked: the later anonymize call still + receives the full original text, so texts above the anonymizer's own + body limit that contain detections keep failing there. + """ + text_chunks: Final = self._split_text_for_analysis( + text=text, + chunk_size_bytes=self.presidio_analyze_chunk_size_bytes, + overlap_chars=PRESIDIO_ANALYZE_CHUNK_OVERLAP_CHARS, + ) + verbose_proxy_logger.debug( + "Presidio analyze: text exceeds %s bytes, analyzing in %s overlapping chunks", + self.presidio_analyze_chunk_size_bytes, + len(text_chunks), + ) + # 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, + ) -> Sequence[PresidioAnalyzeResponseItem] | _PresidioAnonymizeResponse: + async with analyze_semaphore: + return await self.analyze_text( + text=chunk_text, + presidio_config=presidio_config, + request_data=request_data, + ) + + gathered: Final = await asyncio.gather( + *(_analyze_chunk_bounded(chunk_text) for _, chunk_text in text_chunks), + return_exceptions=True, + ) + chunk_results: Final = [] + for result in gathered: + if isinstance(result, BaseException): + raise result + # 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 + # 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: + """ + Validate a configured chunk size, falling back to the default. + + 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 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 + return max(value, 4) + + @staticmethod + def _split_text_for_analysis( + text: str, + chunk_size_bytes: int, + overlap_chars: int, + ) -> Sequence[tuple[int, str]]: + """ + 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 + seen whole by one of the chunks (longer boundary-straddling entities + may be seen only truncated); ``_merge_chunked_analyze_results`` resolves + the duplicate and truncated detections this produces. Returns + ``(char_offset, chunk_text)`` pairs where ``char_offset`` is the + chunk's start position in the original text. + """ + chunks: Final = [] + text_len: Final = len(text) + start = 0 # rebind-ok: chunk cursor advances across the loop + while start < text_len: + # 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] + 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: + break + # Cap the overlap so the next chunk always makes forward progress. + effective_overlap = min(overlap_chars, len(chunk) // 2) + start = max(start + 1, end - effective_overlap) + return chunks + + @staticmethod + def _merge_chunked_analyze_results( + text_chunks: Sequence[tuple[int, str]], + chunk_results: Sequence[Sequence[PresidioAnalyzeResponseItem]], + ) -> list[PresidioAnalyzeResponseItem]: # mutable-ok: analyze_text's declared return type requires list + """ + Remap per-chunk analyzer offsets onto the original text and merge. + + A detection in an overlap region is reported by both neighbouring + chunks, and a boundary entity can additionally be reported truncated by + the chunk that saw only its head or tail. Same-entity-type detections + with overlapping remapped spans are therefore resolved by keeping the + longest span (highest score on ties) — mirroring the same-type conflict + removal Presidio's AnalyzerEngine applies within a single call, and + keeping overlapping spans from corrupting the numbered-token rewriter. + Detections of DIFFERENT entity types may still overlap, exactly as in a + single-call response. The merged list is sorted by position. + """ + remapped: Final = [] + for (char_offset, _), results in zip(text_chunks, chunk_results, strict=True): + for item in results: + item_start = item.get("start") + item_end = item.get("end") + if item_start is not None: + item["start"] = item_start + char_offset + if item_end is not None: + item["end"] = item_end + char_offset + remapped.append(item) + + def _priority(item: PresidioAnalyzeResponseItem) -> tuple[int, float]: + span_start: Final = item.get("start") or 0 + span_end: Final = item.get("end") or 0 + return (-(span_end - span_start), -(item.get("score") or 0.0)) + + merged: Final = [] + kept_spans_by_type: Final = {} + for item in sorted(remapped, key=_priority): + item_start = item.get("start") + item_end = item.get("end") + if item_start is None or item_end is None: + merged.append(item) + continue + kept_spans = kept_spans_by_type.setdefault(str(item.get("entity_type")), []) + if any(item_start < kept_end and kept_start < item_end for kept_start, kept_end in kept_spans): + continue + kept_spans.append((item_start, item_end)) + merged.append(item) + merged.sort(key=lambda r: (r.get("start") or 0, r.get("end") or 0)) + return merged + async def _post_presidio_anonymize( self, text: str, @@ -1392,3 +1631,9 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): 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 + if litellm_params.presidio_analyze_chunk_size_bytes is not None: + # Same validation as __init__: a non-positive value from a guardrail + # update must not silently disable detection via degenerate chunking. + self.presidio_analyze_chunk_size_bytes = self._coerce_analyze_chunk_size( + litellm_params.presidio_analyze_chunk_size_bytes + ) diff --git a/litellm/proxy/guardrails/guardrail_initializers.py b/litellm/proxy/guardrails/guardrail_initializers.py index b377b272e5b..35b6e240d7d 100644 --- a/litellm/proxy/guardrails/guardrail_initializers.py +++ b/litellm/proxy/guardrails/guardrail_initializers.py @@ -104,7 +104,12 @@ def initialize_presidio(litellm_params: LitellmParams, guardrail: Guardrail): 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/litellm/types/guardrails.py b/litellm/types/guardrails.py index 5e21dd2f60c..e48c36de8ba 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -392,6 +392,16 @@ class PresidioConfigModel(PresidioPresidioConfigModelUserInterface): default=None, description="Path to a JSON file containing ad-hoc recognizers for Presidio", ) + presidio_analyze_chunk_size_bytes: int | None = Field( + default=None, + description=( + "Maximum UTF-8 bytes of text sent in a single Presidio /analyze call. " + "Longer texts are split into overlapping chunks of at most this size " + "and the merged results are remapped onto the original text. " + "Defaults to 500000; set it below your analyzer deployment's request " + "body limit, leaving headroom for the rest of the analyze payload." + ), + ) mock_redacted_text: dict | None = Field(default=None, description="Mock redacted text for testing") 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 acb43bc5b74..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()) @@ -2893,3 +2824,559 @@ async def test_stream_pii_unmasking_passthrough_when_no_tokens(mock_user_api_key chunks.append(chunk) assert chunks == [raw_chunk] + + +# --------------------------------------------------------------------------- +# Chunked /analyze tests (LIT-4785) +# Oversized texts must be split into overlapping chunks before /analyze, with +# per-chunk offsets remapped onto the original text. +# --------------------------------------------------------------------------- + +CHUNK_MARKER_ONE = "4111-0001" +CHUNK_MARKER_TWO = "4111-0002" + + +def _make_marker_session_iterator( + recorded_analyze_payloads, + analyzer_body_limit_bytes=None, + recorded_anonymize_payloads=None, +): + """Mock session behaving like a real Presidio pair. + + /analyze returns a CREDIT_CARD detection for every ``4111-NNNN`` marker in + the posted text (chunk-local offsets, like the real analyzer). When + ``analyzer_body_limit_bytes`` is set, oversized /analyze bodies get the + HTTP 413 from LIT-4785. /anonymize replaces the given spans in the posted + text. + """ + import json as json_module + import re as re_module + + @asynccontextmanager + async def mock_iterator(): + class MockResponse: + def __init__(self, status, body): + self.status = status + self.content_type = "application/json" + self.headers = {"Content-Type": "application/json"} + self._body = body + + async def text(self): + 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): + payload = json + 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: + return MockResponse( + 413, + { + "error": "Request body too large. /analyze accepts at most " + f"{analyzer_body_limit_bytes} bytes; larger documents must be " + "chunked by the caller." + }, + ) + results = [ + { + "entity_type": "CREDIT_CARD", + "start": m.start(), + "end": m.end(), + "score": 1.0, + } + for m in re_module.finditer(r"4111-\d{4}", text) + ] + return MockResponse(200, results) + 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) + for r in items: + text = text[: r["start"]] + "<" + r["entity_type"] + ">" + text[r["end"] :] + return MockResponse( + 200, + { + "text": text, + "items": [{"entity_type": r["entity_type"]} for r in items], + }, + ) + + async def __aenter__(self): + return self + + async def __aexit__(self, *args): + return False + + yield MockSession() + + return mock_iterator + + +def _chunking_guardrail(chunk_size_bytes=100, **kwargs): + return _OPTIONAL_PresidioPIIMasking( + presidio_analyzer_api_base="http://test-analyzer/", + presidio_anonymizer_api_base="http://test-anonymizer/", + presidio_analyze_chunk_size_bytes=chunk_size_bytes, + mock_testing=False, + **kwargs, + ) + + +def _oversized_marker_text(): + """~258-char text with markers in the 1st and 3rd 100-byte chunk.""" + filler = "x" * 60 + return filler + CHUNK_MARKER_ONE + filler + filler + CHUNK_MARKER_TWO + filler + + +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) + assert len(chunks) > 1 + for offset, chunk in chunks: + assert len(chunk.encode("utf-8")) <= 100 + assert text[offset : offset + len(chunk)] == chunk + assert chunks[0][0] == 0 + assert chunks[-1][0] + len(chunks[-1][1]) == len(text) + for (prev_off, prev_chunk), (next_off, _) in zip(chunks, chunks[1:]): + # consecutive chunks overlap (or at least touch) and make progress + assert next_off <= prev_off + len(prev_chunk) + assert next_off > prev_off + + +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) + assert len(chunks) > 1 + for offset, chunk in chunks: + assert len(chunk.encode("utf-8")) <= 64 + assert text[offset : offset + len(chunk)] == chunk + assert chunks[-1][0] + len(chunks[-1][1]) == len(text) + + +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) + assert chunks == [(0, text)] + + +@pytest.mark.asyncio +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={}) + assert len(payloads) == 1 + assert payloads[0]["text"] == text + assert len(results) == 1 + assert text[results[0]["start"] : results[0]["end"]] == CHUNK_MARKER_ONE + + +@pytest.mark.asyncio +async def test_analyze_text_chunks_oversized_text_and_remaps_offsets(): + """Regression test for LIT-4785. + + The mock analyzer rejects bodies over 100 bytes with HTTP 413 (like the + reporter's deployment): on unfixed code the single oversized /analyze call + fails closed; with chunking every call stays under the limit and the + detections come back with offsets remapped onto the original text. + The duplicate detection from the overlap region must be deduplicated. + """ + guardrail = _chunking_guardrail( + chunk_size_bytes=100, + pii_entities_config={"CREDIT_CARD": PiiAction.MASK}, + ) + payloads = [] + text = _oversized_marker_text() + with patch.object( + guardrail, + "_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={}) + assert len(payloads) > 1 + for payload in payloads: + assert len(payload["text"].encode("utf-8")) <= 100 + assert [text[r["start"] : r["end"]] for r in results] == [ + CHUNK_MARKER_ONE, + CHUNK_MARKER_TWO, + ] + + +@pytest.mark.asyncio +async def test_check_pii_masks_oversized_text_with_chunking(): + guardrail = _chunking_guardrail( + chunk_size_bytes=100, + pii_entities_config={"CREDIT_CARD": PiiAction.MASK}, + ) + analyze_payloads = [] + anonymize_payloads = [] + text = _oversized_marker_text() + with patch.object( + guardrail, + "_get_session_iterator", + _make_marker_session_iterator( + analyze_payloads, + analyzer_body_limit_bytes=100, + recorded_anonymize_payloads=anonymize_payloads, + ), + ): + 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 + # anonymize still receives the full text with globally remapped offsets + assert len(anonymize_payloads) == 1 + assert anonymize_payloads[0]["text"] == text + + +@pytest.mark.asyncio +async def test_output_parse_pii_numbered_tokens_across_chunks(): + """Numbered tokens slice the ORIGINAL text at the remapped offsets; a + chunk-local offset would store the wrong substring in pii_tokens and + corrupt the later unmask.""" + guardrail = _chunking_guardrail( + chunk_size_bytes=100, + pii_entities_config={"CREDIT_CARD": PiiAction.MASK}, + output_parse_pii=True, + ) + payloads = [] + request_data = {} + text = _oversized_marker_text() + with patch.object( + guardrail, + "_get_session_iterator", + _make_marker_session_iterator(payloads, analyzer_body_limit_bytes=100), + ): + masked = await guardrail.check_pii( + text=text, + output_parse_pii=True, + presidio_config=None, + request_data=request_data, + ) + assert masked.count("") == 1 + assert masked.count("") == 1 + pii_tokens = request_data["metadata"]["pii_tokens"] + assert pii_tokens[""] == CHUNK_MARKER_ONE + assert pii_tokens[""] == CHUNK_MARKER_TWO + + +@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 + failing /analyze call (fail closed when PII protection is configured).""" + guardrail = _chunking_guardrail( + chunk_size_bytes=100, + pii_entities_config={"CREDIT_CARD": PiiAction.MASK}, + ) + payloads = [] + text = _oversized_marker_text() + with patch.object( + guardrail, + "_get_session_iterator", + # every chunk is rejected: limit below the chunk size + _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={}) + + +def test_presidio_analyze_chunk_size_default_and_validation(): + from litellm.constants import DEFAULT_PRESIDIO_ANALYZE_CHUNK_SIZE_BYTES + + 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) + 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) + assert custom.presidio_analyze_chunk_size_bytes == 1234 + + +def test_update_in_memory_applies_analyze_chunk_size(): + guardrail = _OPTIONAL_PresidioPIIMasking(mock_testing=True) + params = LitellmParams( + guardrail="presidio", + mode="pre_call", + presidio_analyze_chunk_size_bytes=99_000, + ) + guardrail.update_in_memory_litellm_params(params) + assert guardrail.presidio_analyze_chunk_size_bytes == 99_000 + + +def test_merge_drops_truncated_same_type_fragment_from_overlap(): + """A boundary entity seen truncated by chunk 1 and whole by chunk 2 must + merge to the single full span; keeping both overlapping spans corrupts the + numbered-token rewriter and double-counts entities.""" + truncated = {"entity_type": "IP_ADDRESS", "start": 10, "end": 21, "score": 0.6} + full_local = {"entity_type": "IP_ADDRESS", "start": 5, "end": 18, "score": 0.95} + merged = _OPTIONAL_PresidioPIIMasking._merge_chunked_analyze_results( + text_chunks=[(0, "x" * 21), (5, "x" * 25)], + chunk_results=[[truncated], [full_local]], + ) + assert len(merged) == 1 + assert (merged[0]["start"], merged[0]["end"]) == (10, 23) + assert merged[0]["score"] == 0.95 + + +def test_merge_exact_duplicate_keeps_higher_score(): + low = {"entity_type": "EMAIL_ADDRESS", "start": 3, "end": 9, "score": 0.4} + high = {"entity_type": "EMAIL_ADDRESS", "start": 0, "end": 6, "score": 0.9} + merged = _OPTIONAL_PresidioPIIMasking._merge_chunked_analyze_results( + text_chunks=[(0, "x" * 9), (3, "x" * 9)], + chunk_results=[[low], [high]], + ) + assert len(merged) == 1 + assert merged[0]["score"] == 0.9 + + +def test_merge_preserves_cross_type_overlap(): + """Single-call Presidio returns overlapping detections of DIFFERENT types + (e.g. URL inside EMAIL_ADDRESS); the chunk merge must not drop those.""" + email = {"entity_type": "EMAIL_ADDRESS", "start": 0, "end": 20, "score": 1.0} + url = {"entity_type": "URL", "start": 5, "end": 20, "score": 0.5} + merged = _OPTIONAL_PresidioPIIMasking._merge_chunked_analyze_results( + text_chunks=[(0, "x" * 25)], + chunk_results=[[email, url]], + ) + assert len(merged) == 2 + + +def test_update_in_memory_coerces_invalid_chunk_size(): + from litellm.constants import DEFAULT_PRESIDIO_ANALYZE_CHUNK_SIZE_BYTES + + guardrail = _OPTIONAL_PresidioPIIMasking(mock_testing=True, presidio_analyze_chunk_size_bytes=99_000) + params = LitellmParams( + guardrail="presidio", + mode="pre_call", + presidio_analyze_chunk_size_bytes=-1, + ) + guardrail.update_in_memory_litellm_params(params) + assert guardrail.presidio_analyze_chunk_size_bytes == DEFAULT_PRESIDIO_ANALYZE_CHUNK_SIZE_BYTES + + +def test_split_text_handles_chunk_size_below_char_width(): + chunks = _OPTIONAL_PresidioPIIMasking._split_text_for_analysis( + text="\U0001f642\U0001f642", chunk_size_bytes=3, overlap_chars=8 + ) + assert all(chunk for _, chunk in chunks) + assert chunks[-1][0] + len(chunks[-1][1]) == 2 + + +@pytest.mark.asyncio +async def test_tiny_chunk_size_with_multibyte_text_terminates(): + """chunk_size below one character's UTF-8 width must not recurse forever; + the constructor floors the value to the widest character width.""" + 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)): + results = await guardrail.analyze_text( + text="\U0001f642\U0001f642\U0001f642ab", presidio_config=None, request_data={} + ) + assert results == [] + assert len(payloads) >= 2 + + +@pytest.mark.asyncio +async def test_chunked_analyze_concurrency_is_bounded(): + 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 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 diff --git a/tests/test_litellm/proxy/guardrails/test_init_guardrails.py b/tests/test_litellm/proxy/guardrails/test_init_guardrails.py index 82363302d2e..4eed1aa509f 100644 --- a/tests/test_litellm/proxy/guardrails/test_init_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/test_init_guardrails.py @@ -118,3 +118,38 @@ def test_initialize_guardrail_sets_run_in_parallel(config_value, expected): custom_guardrail = guardrail_handler.guardrail_id_to_custom_guardrail[result["guardrail_id"]] assert custom_guardrail.run_in_parallel is expected + + +def test_initialize_presidio_forwards_analyze_chunk_size_bytes(): + """Regression (LIT-4785): `presidio_analyze_chunk_size_bytes` set in + config.yaml must reach the guardrail instance. The field lives on + PresidioConfigModel, so LitellmParams parses it, but initialize_presidio + enumerates its constructor kwargs explicitly and would silently drop it. + """ + import litellm + from litellm.proxy.guardrails.guardrail_hooks.presidio import ( + _OPTIONAL_PresidioPIIMasking, + ) + + test_guardrail = { + "guardrail_name": "test_presidio_chunk_size", + "litellm_params": { + "guardrail": SupportedGuardrailIntegrations.PRESIDIO.value, + "mode": "pre_call", + "presidio_analyzer_api_base": "https://fakelink.com/v1/presidio/analyze", + "presidio_anonymizer_api_base": "https://fakelink.com/v1/presidio/anonymize", + "presidio_analyze_chunk_size_bytes": 250_000, + }, + } + + guardrail_handler = InMemoryGuardrailHandler() + guardrail_handler.initialize_guardrail(guardrail=test_guardrail) + + initialized = [ + callback + for callback in litellm.callbacks + if isinstance(callback, _OPTIONAL_PresidioPIIMasking) + and callback.guardrail_name == "test_presidio_chunk_size" + ] + assert initialized, "presidio guardrail was not registered as a callback" + assert initialized[-1].presidio_analyze_chunk_size_bytes == 250_000