mirror of
https://github.com/BerriAI/litellm.git
synced 2026-08-28 05:25:59 +00:00
fix(presidio): chunk oversized text before /analyze so large content blocks do not fail (#38483)
* fix(presidio): chunk oversized text before /analyze so large content blocks do not fail The Presidio PII guardrail sent each content block to the analyzer as a single /analyze call with no size check. Analyzer deployments commonly cap the request body (the reporting deployment rejects bodies over 1,000,000 bytes with HTTP 413), so large blocks failed closed, and analyzer latency grew linearly with payload size. analyze_text now splits texts larger than presidio_analyze_chunk_size_bytes (default 500,000 UTF-8 bytes, configurable per guardrail) into overlapping chunks, analyzes them concurrently, remaps each detection's start/end onto the original text, and deduplicates detections from the overlap regions. Anonymization, blocked-entity checks, score filtering, numbered-token unmasking, telemetry, and the dashboard entity positions all consume the remapped global offsets unchanged. Resolves LIT-4785 Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * fix(presidio): review-round hardening for chunked analyze - measure the chunk budget on the JSON-serialized text (non-ASCII escapes expand beyond raw UTF-8, so a raw-byte budget could still exceed the analyzer body limit) - share the chunk fan-out semaphore per event loop and instance instead of per call, so many oversized blocks cannot multiply concurrent analyzer calls - apply configured score thresholds and deny list per chunk BEFORE overlap resolution, so a below-threshold span cannot displace a detection the thresholds keep Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> --------- Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
yucheng-berri
GitHub
parent
1ab6fd89d2
commit
239ec955dc
6 changed files with 893 additions and 108 deletions
|
|
@ -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)))
|
||||
|
|
|
|||
|
|
@ -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
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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 <PERSON>, my email is <EMAIL_ADDRESS>, phone <PHONE_NUMBER>"
|
||||
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, {"<PERSON_1>": "Bobby"}
|
||||
)
|
||||
result = _OPTIONAL_PresidioPIIMasking._unmask_sse_bytes_chunk(chunk, {"<PERSON_1>": "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, {"<PERSON_1>": "Bobby"}
|
||||
)
|
||||
result = _OPTIONAL_PresidioPIIMasking._unmask_sse_bytes_chunk(chunk, {"<PERSON_1>": "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("<CREDIT_CARD>") == 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("<CREDIT_CARD_1>") == 1
|
||||
assert masked.count("<CREDIT_CARD_2>") == 1
|
||||
pii_tokens = request_data["metadata"]["pii_tokens"]
|
||||
assert pii_tokens["<CREDIT_CARD_1>"] == CHUNK_MARKER_ONE
|
||||
assert pii_tokens["<CREDIT_CARD_2>"] == 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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue