scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-07 08:26:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

Merge pull request #22627 from BerriAI/litellm_presidio_antropic_edge_case

Browse source 
fix: presidio edge case with antropic handle on pii token leak
This commit is contained in:
Harshit Jain 2026-03-10 21:16:56 +05:30 • committed by GitHub
commit 424a96580b
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 872 additions and 150 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

174
litellm/proxy/guardrails/guardrail_hooks/presidio.py
View file

@ -36,7 +36,6 @@ from litellm.types.utils import GenericGuardrailAPIInputs
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm._uuid import uuid
from litellm.caching.caching import DualCache
from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException
from litellm.integrations.custom_guardrail import (
@ -100,6 +99,16 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
self.mock_redacted_text = mock_redacted_text
self.output_parse_pii = output_parse_pii or False
self.apply_to_output = apply_to_output
# When output_parse_pii or apply_to_output is enabled, the guardrail must
# also run on post_call to unmask/mask the response. Expand the event_hook
# so should_run_guardrail returns True for both pre_call and post_call.
if (self.output_parse_pii or self.apply_to_output) and not logging_only:
current_hook = self.event_hook
if isinstance(current_hook, str) and current_hook != "post_call":
self.event_hook = [current_hook, "post_call"]
elif isinstance(current_hook, list) and "post_call" not in current_hook:
self.event_hook = current_hook + ["post_call"]
self.pii_entities_config: Dict[Union[PiiEntityType, str], PiiAction] = (
pii_entities_config or {}
)
@ -475,13 +484,15 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
new_text = text
if redacted_text is not None:
verbose_proxy_logger.debug("redacted_text: %s", redacted_text)
for item in redacted_text["items"]:
# Process items in reverse order by start position so that
# replacing later spans first does not shift earlier coordinates.
for item in sorted(
redacted_text["items"], key=lambda x: x["start"], reverse=True
):
start = item["start"]
end = item["end"]
replacement = item["text"] # replacement token
if item["operator"] == "replace" and output_parse_pii is True:
# check if token in dict
# if exists, add a uuid to the replacement token for swapping back to the original text in llm response output parsing
if request_data is None:
verbose_proxy_logger.warning(
"Presidio anonymize_text called without request_data — "
@ -489,17 +500,28 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
"This may indicate a missing caller update."
)
request_data = {}
if "pii_tokens" not in request_data:
request_data["pii_tokens"] = {}
pii_tokens = request_data["pii_tokens"]
# Store pii_tokens in metadata to avoid leaking to LLM providers.
# Providers like Anthropic reject unknown top-level fields.
if not request_data.get("metadata"):
request_data["metadata"] = {}
if "pii_tokens" not in request_data["metadata"]:
request_data["metadata"]["pii_tokens"] = {}
pii_tokens = request_data["metadata"]["pii_tokens"]
# Always append a UUID to ensure the replacement token is unique to this request and session.
# This prevents collisions where the LLM might hallucinate a generic token like [PHONE_NUMBER].
replacement = f"{replacement}_{str(uuid.uuid4())[:12]}"
# Append a sequential number to make each token unique
# per request, so unmasking maps back to the correct
# original value. Format: <PHONE_NUMBER_1>, <PHONE_NUMBER_2>
# This is LLM-friendly and degrades gracefully if the
# LLM doesn't echo the token verbatim.
seq = len(pii_tokens) + 1
if replacement.endswith(">"):
replacement = f"{replacement[:-1]}_{seq}>"
else:
replacement = f"{replacement}_{seq}"
pii_tokens[replacement] = new_text[
start:end
] # get text it'll replace
# Use ORIGINAL text (not new_text) since start/end
# reference the original text's coordinates.
pii_tokens[replacement] = text[start:end]
new_text = new_text[:start] + replacement + new_text[end:]
entity_type = item.get("entity_type", None)
@ -507,12 +529,13 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
masked_entity_count[entity_type] = (
masked_entity_count.get(entity_type, 0) + 1
)
# When output_parse_pii is True, new_text contains UUID-suffixed
# tokens that match the keys in pii_tokens. Returning
# redacted_text["text"] (Presidio's original output) would send
# un-suffixed tokens to the LLM, making unmasking impossible.
# When output_parse_pii is True, new_text contains sequentially
# numbered tokens (e.g. <PHONE_NUMBER_1>) that match the keys
# in pii_tokens. Returning redacted_text["text"] (Presidio's
# original output) would send un-numbered tokens to the LLM,
# making unmasking impossible.
# When output_parse_pii is False, new_text == redacted_text["text"]
# because no UUID suffix is appended.
# because no suffix is appended.
return new_text
else:
raise Exception("Invalid anonymizer response: received None")
@ -544,8 +567,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
filtered_results: List[PresidioAnalyzeResponseItem] = []
deny_list_strings = [
getattr(x, "value", str(x))
for x in self.presidio_entities_deny_list
getattr(x, "value", str(x)) for x in self.presidio_entities_deny_list
]
for item in analyze_results:
entity_type = item.get("entity_type")
@ -884,6 +906,10 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
)
if self.apply_to_output is True:
if self._is_anthropic_message_response(response):
return await self._process_anthropic_response_for_pii(
response=response, request_data=data, mode="mask"
)
return await self._mask_output_response(
response=response, request_data=data
)
@ -899,6 +925,10 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
request_data=data,
mode="unmask",
)
elif self._is_anthropic_message_response(response):
await self._process_anthropic_response_for_pii(
response=response, request_data=data, mode="unmask"
)
return response
@staticmethod
@ -927,6 +957,57 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
break
return text
@staticmethod
def _is_anthropic_message_response(response: Any) -> bool:
"""Check if the response is an Anthropic native message dict."""
return (
isinstance(response, dict)
and response.get("type") == "message"
and isinstance(response.get("content"), list)
)
async def _process_anthropic_response_for_pii(
self,
response: dict,
request_data: dict,
mode: Literal["mask", "unmask"],
) -> dict:
"""
Process an Anthropic native message dict for PII masking/unmasking.
Handles content blocks with type == "text".
"""
metadata = (request_data.get("metadata") or {}) if request_data else {}
pii_tokens = metadata.get("pii_tokens", {})
if not pii_tokens and mode == "unmask":
verbose_proxy_logger.debug(
"No pii_tokens in metadata for Anthropic response unmask"
)
presidio_config = self.get_presidio_settings_from_request_data(
request_data or {}
)
content = response.get("content")
if not isinstance(content, list):
return response
for block in content:
if not isinstance(block, dict) or block.get("type") != "text":
continue
text_value = block.get("text")
if text_value is None:
continue
if mode == "unmask":
block["text"] = self._unmask_pii_text(text_value, pii_tokens)
elif mode == "mask":
block["text"] = await self.check_pii(
text=text_value,
output_parse_pii=False,
presidio_config=presidio_config,
request_data=request_data,
)
return response
async def _process_response_for_pii(
self,
response: ModelResponse,
@ -937,10 +1018,11 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
Helper to recursively process a ModelResponse for PII.
Handles all choices and tool calls.
"""
pii_tokens = request_data.get("pii_tokens", {}) if request_data else {}
metadata = (request_data.get("metadata") or {}) if request_data else {}
pii_tokens = metadata.get("pii_tokens", {})
if not pii_tokens and mode == "unmask":
verbose_proxy_logger.debug(
"No pii_tokens found in request_data — nothing to unmask"
"No pii_tokens found in request_data['metadata'] — nothing to unmask"
)
presidio_config = self.get_presidio_settings_from_request_data(
request_data or {}
@ -1045,7 +1127,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
user_api_key_dict: UserAPIKeyAuth,
response: Any,
request_data: dict,
) -> AsyncGenerator[ModelResponseStream, None]:
) -> AsyncGenerator[Union[ModelResponseStream, bytes], None]:
"""
Process streaming response chunks to unmask PII tokens when needed.
"""
@ -1062,8 +1144,20 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
async for chunk in response:
if isinstance(chunk, ModelResponseStream):
all_chunks.append(chunk)
elif isinstance(chunk, bytes):
# Anthropic native SSE: pass through as-is
yield chunk # type: ignore[misc]
continue
if not all_chunks:
# All chunks were Anthropic native SSE bytes — output
# masking cannot be applied to raw bytes. Log a warning
# so operators know PII masking was skipped for this stream.
verbose_proxy_logger.warning(
"Presidio apply_to_output: streaming response contained only "
"bytes chunks (Anthropic native SSE). Output PII masking was "
"skipped for this response."
)
return
assembled_model_response = stream_chunk_builder(
@ -1099,10 +1193,11 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
return
# --- PII unmasking path (output_parse_pii=True) ---
pii_tokens = request_data.get("pii_tokens", {}) if request_data else {}
metadata = (request_data.get("metadata") or {}) if request_data else {}
pii_tokens = metadata.get("pii_tokens", {})
if not pii_tokens and request_data:
verbose_proxy_logger.debug(
"No pii_tokens in request_data for streaming unmask path"
"No pii_tokens in request_data['metadata'] for streaming unmask path"
)
if not (self.output_parse_pii and pii_tokens):
async for chunk in response:
@ -1114,6 +1209,10 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
async for chunk in response:
if isinstance(chunk, ModelResponseStream):
remaining_chunks.append(chunk)
elif isinstance(chunk, bytes):
# Anthropic native SSE: pass through as-is
yield chunk # type: ignore[misc]
continue
if not remaining_chunks:
return
@ -1191,15 +1290,24 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
"""
texts = inputs.get("texts", [])
# When input_type is "response" and pii_tokens are available,
# unmask the text instead of masking it.
metadata = (request_data.get("metadata") or {}) if request_data else {}
pii_tokens = metadata.get("pii_tokens", {})
new_texts = []
for text in texts:
modified_text = await self.check_pii(
text=text,
output_parse_pii=self.output_parse_pii,
presidio_config=None,
request_data=request_data or {},
)
new_texts.append(modified_text)
if input_type == "response" and pii_tokens:
for text in texts:
new_texts.append(self._unmask_pii_text(text, pii_tokens))
else:
for text in texts:
modified_text = await self.check_pii(
text=text,
output_parse_pii=self.output_parse_pii,
presidio_config=None,
request_data=request_data or {},
)
new_texts.append(modified_text)
inputs["texts"] = new_texts
return inputs

226
tests/guardrails_tests/test_presidio_pii.py
View file

@ -1,20 +1,19 @@
import sys
import os
import io, asyncio
import pytest
import time
from litellm import mock_completion
from unittest.mock import MagicMock, AsyncMock, patch
from unittest.mock import patch
sys.path.insert(0, os.path.abspath("../.."))
import litellm
from litellm.proxy.guardrails.guardrail_hooks.presidio import _OPTIONAL_PresidioPIIMasking, PresidioPerRequestConfig
from litellm.proxy.guardrails.guardrail_hooks.presidio import (
_OPTIONAL_PresidioPIIMasking,
PresidioPerRequestConfig,
)
from litellm.types.guardrails import PiiEntityType, PiiAction
from litellm.proxy._types import UserAPIKeyAuth
from litellm.caching.caching import DualCache
from litellm.exceptions import BlockedPiiEntityError
from litellm.types.utils import CallTypes as LitellmCallTypes
@pytest.mark.asyncio
@ -26,42 +25,37 @@ async def test_presidio_with_entities_config():
PiiEntityType.CREDIT_CARD: PiiAction.MASK,
PiiEntityType.EMAIL_ADDRESS: PiiAction.MASK,
}
presidio_guardrail = _OPTIONAL_PresidioPIIMasking(
pii_entities_config=pii_entities_config,
presidio_analyzer_api_base=os.environ.get("PRESIDIO_ANALYZER_API_BASE"),
presidio_anonymizer_api_base=os.environ.get("PRESIDIO_ANONYMIZER_API_BASE")
presidio_anonymizer_api_base=os.environ.get("PRESIDIO_ANONYMIZER_API_BASE"),
)
# Test text with different PII types
test_text = "My credit card number is 4111-1111-1111-1111, my email is test@example.com, and my phone is 555-123-4567"
# Test the analyze request configuration
analyze_request = presidio_guardrail._get_presidio_analyze_request_payload(
text=test_text,
presidio_config=None,
request_data={}
text=test_text, presidio_config=None, request_data={}
)
# Verify entities were passed correctly
assert "entities" in analyze_request
assert set(analyze_request["entities"]) == set(pii_entities_config.keys())
# Test the check_pii method - this will call the actual Presidio API
redacted_text = await presidio_guardrail.check_pii(
text=test_text,
output_parse_pii=True,
presidio_config=None,
request_data={}
text=test_text, output_parse_pii=True, presidio_config=None, request_data={}
)
# Verify PII has been masked/replaced/redacted in the result
assert "4111-1111-1111-1111" not in redacted_text
assert "test@example.com" not in redacted_text
# Since this entity is not in the config, it should not be masked
assert "555-123-4567" in redacted_text
# The specific replacements will vary based on Presidio's implementation
print(f"Redacted text: {redacted_text}")
@ -73,10 +67,12 @@ async def test_presidio_apply_guardrail():
presidio_guardrail = _OPTIONAL_PresidioPIIMasking(
pii_entities_config={},
presidio_analyzer_api_base=os.environ.get("PRESIDIO_ANALYZER_API_BASE"),
presidio_anonymizer_api_base=os.environ.get("PRESIDIO_ANONYMIZER_API_BASE")
presidio_anonymizer_api_base=os.environ.get("PRESIDIO_ANONYMIZER_API_BASE"),
)
test_text = "My credit card number is 4111-1111-1111-1111 and my email is test@example.com"
test_text = (
"My credit card number is 4111-1111-1111-1111 and my email is test@example.com"
)
response = await presidio_guardrail.apply_guardrail(
inputs={"texts": [test_text]},
request_data={},
@ -91,6 +87,7 @@ async def test_presidio_apply_guardrail():
assert "4111-1111-1111-1111" not in modified_text
assert "test@example.com" not in modified_text
@pytest.mark.asyncio
async def test_presidio_with_blocked_entities():
"""Test for Presidio guardrail with blocked entities - requires actual Presidio API"""
@ -100,36 +97,33 @@ async def test_presidio_with_blocked_entities():
PiiEntityType.CREDIT_CARD: PiiAction.BLOCK, # This entity should cause a block
PiiEntityType.EMAIL_ADDRESS: PiiAction.MASK, # This entity should be masked
}
presidio_guardrail = _OPTIONAL_PresidioPIIMasking(
pii_entities_config=pii_entities_config,
presidio_analyzer_api_base=os.environ.get("PRESIDIO_ANALYZER_API_BASE"),
presidio_anonymizer_api_base=os.environ.get("PRESIDIO_ANONYMIZER_API_BASE")
presidio_anonymizer_api_base=os.environ.get("PRESIDIO_ANONYMIZER_API_BASE"),
)
# Test text with blocked PII type
test_text = "My credit card number is 4111-1111-1111-1111 and my email is test@example.com"
test_text = (
"My credit card number is 4111-1111-1111-1111 and my email is test@example.com"
)
# Verify the analyze request configuration
analyze_request = presidio_guardrail._get_presidio_analyze_request_payload(
text=test_text,
presidio_config=None,
request_data={}
text=test_text, presidio_config=None, request_data={}
)
# Verify entities were passed correctly
assert "entities" in analyze_request
assert set(analyze_request["entities"]) == set(pii_entities_config.keys())
# Test that BlockedPiiEntityError is raised when check_pii is called
with pytest.raises(BlockedPiiEntityError) as excinfo:
await presidio_guardrail.check_pii(
text=test_text,
output_parse_pii=True,
presidio_config=None,
request_data={}
text=test_text, output_parse_pii=True, presidio_config=None, request_data={}
)
# Verify the error contains the correct entity type
assert excinfo.value.entity_type == PiiEntityType.CREDIT_CARD
assert excinfo.value.guardrail_name == presidio_guardrail.guardrail_name
@ -143,37 +137,40 @@ async def test_presidio_pre_call_hook_with_blocked_entities():
PiiEntityType.CREDIT_CARD: PiiAction.BLOCK, # This entity should cause a block
PiiEntityType.EMAIL_ADDRESS: PiiAction.MASK, # This entity should be masked
}
presidio_guardrail = _OPTIONAL_PresidioPIIMasking(
pii_entities_config=pii_entities_config,
presidio_analyzer_api_base=os.environ.get("PRESIDIO_ANALYZER_API_BASE"),
presidio_anonymizer_api_base=os.environ.get("PRESIDIO_ANONYMIZER_API_BASE")
presidio_anonymizer_api_base=os.environ.get("PRESIDIO_ANONYMIZER_API_BASE"),
)
# Create a sample chat completion request with PII data
data = {
"messages": [
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": "My credit card is 4111-1111-1111-1111 and my email is test@example.com."}
{
"role": "user",
"content": "My credit card is 4111-1111-1111-1111 and my email is test@example.com.",
},
],
"model": "gpt-3.5-turbo"
"model": "gpt-3.5-turbo",
}
# Mock objects needed for the pre-call hook
user_api_key_dict = UserAPIKeyAuth(api_key="test_key")
cache = DualCache()
# Call the pre-call hook and expect BlockedPiiEntityError
with pytest.raises(BlockedPiiEntityError) as excinfo:
await presidio_guardrail.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=cache,
data=data,
call_type="completion"
call_type="completion",
)
print(f"got error: {excinfo}")
# Verify the error contains the correct entity type
assert excinfo.value.entity_type == PiiEntityType.CREDIT_CARD
assert excinfo.value.guardrail_name == presidio_guardrail.guardrail_name
@ -188,44 +185,46 @@ async def test_presidio_pre_call_hook_with_different_call_types(call_type):
PiiEntityType.CREDIT_CARD: PiiAction.MASK,
PiiEntityType.EMAIL_ADDRESS: PiiAction.MASK,
}
presidio_guardrail = _OPTIONAL_PresidioPIIMasking(
pii_entities_config=pii_entities_config,
presidio_analyzer_api_base=os.environ.get("PRESIDIO_ANALYZER_API_BASE"),
presidio_anonymizer_api_base=os.environ.get("PRESIDIO_ANONYMIZER_API_BASE")
presidio_anonymizer_api_base=os.environ.get("PRESIDIO_ANONYMIZER_API_BASE"),
)
# Create a sample request with PII data
data = {
"messages": [
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": "My credit card is 4111-1111-1111-1111 and my email is test@example.com. My phone number is 555-123-4567"}
{
"role": "user",
"content": "My credit card is 4111-1111-1111-1111 and my email is test@example.com. My phone number is 555-123-4567",
},
],
"model": "gpt-3.5-turbo"
"model": "gpt-3.5-turbo",
}
# Mock objects needed for the pre-call hook
user_api_key_dict = UserAPIKeyAuth(api_key="test_key")
cache = DualCache()
# Call the pre-call hook with the specified call type
modified_data = await presidio_guardrail.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=cache,
data=data,
call_type=call_type
user_api_key_dict=user_api_key_dict, cache=cache, data=data, call_type=call_type
)
# Verify the messages have been modified to mask PII
assert modified_data["messages"][0]["content"] == "You are a helpful assistant." # System prompt should be unchanged
assert (
modified_data["messages"][0]["content"] == "You are a helpful assistant."
) # System prompt should be unchanged
user_message = modified_data["messages"][1]["content"]
assert "4111-1111-1111-1111" not in user_message
assert "test@example.com" not in user_message
# Since this entity is not in the config, it should not be masked
assert "555-123-4567" in user_message
print(f"Modified user message for call_type={call_type}: {user_message}")
@ -243,7 +242,7 @@ def test_validate_environment_missing_http(base_url):
# Use patch.dict to temporarily modify environment variables only for this test
env_vars = {
"PRESIDIO_ANALYZER_API_BASE": f"{base_url}/analyze",
"PRESIDIO_ANONYMIZER_API_BASE": f"{base_url}/anonymize"
"PRESIDIO_ANONYMIZER_API_BASE": f"{base_url}/anonymize",
}
with patch.dict(os.environ, env_vars):
pii_masking.validate_environment()
@ -294,8 +293,12 @@ async def test_output_parsing():
new_response = await pii_masking.async_post_call_success_hook(
user_api_key_dict=UserAPIKeyAuth(),
data={
"messages": [{"role": "system", "content": "You are an helpfull assistant"}],
"pii_tokens": {"<PERSON>": "Jane Doe", "<PHONE_NUMBER>": "034453334"},
"messages": [
{"role": "system", "content": "You are an helpfull assistant"}
],
"metadata": {
"pii_tokens": {"<PERSON>": "Jane Doe", "<PHONE_NUMBER>": "034453334"}
},
},
response=response,
)
@ -440,24 +443,26 @@ async def test_presidio_pii_masking_logging_output_only_no_pre_api_hook():
@pytest.mark.asyncio
@patch.dict(os.environ, {
"PRESIDIO_ANALYZER_API_BASE": "http://localhost:5002",
"PRESIDIO_ANONYMIZER_API_BASE": "http://localhost:5001"
})
@patch.dict(
os.environ,
{
"PRESIDIO_ANALYZER_API_BASE": "http://localhost:5002",
"PRESIDIO_ANONYMIZER_API_BASE": "http://localhost:5001",
},
)
async def test_presidio_pii_masking_logging_output_only_logged_response_guardrails_config():
from typing import Dict, List, Optional
import litellm
from litellm.proxy.guardrails.init_guardrails import initialize_guardrails
from litellm.types.guardrails import (
GuardrailItem,
GuardrailItemSpec,
GuardrailEventHooks,
)
litellm.set_verbose = True
# Environment variables are now patched via the decorator instead of setting them directly
guardrails_config: List[Dict[str, GuardrailItemSpec]] = [
{
"pii_masking": {
@ -499,60 +504,53 @@ async def test_presidio_pii_masking_logging_output_only_logged_response_guardrai
async def test_presidio_language_configuration():
"""Test that presidio_language parameter is properly set and used in analyze requests"""
litellm._turn_on_debug()
# Test with German language using mock testing to avoid API calls
presidio_guardrail_de = _OPTIONAL_PresidioPIIMasking(
pii_entities_config={},
presidio_language="de",
mock_testing=True # This bypasses the API validation
mock_testing=True, # This bypasses the API validation
)
test_text = "Meine Telefonnummer ist +49 30 12345678"
# Test the analyze request configuration
analyze_request = presidio_guardrail_de._get_presidio_analyze_request_payload(
text=test_text,
presidio_config=None,
request_data={}
text=test_text, presidio_config=None, request_data={}
)
# Verify the language is set to German
assert analyze_request["language"] == "de"
assert analyze_request["text"] == test_text
# Test with Spanish language
presidio_guardrail_es = _OPTIONAL_PresidioPIIMasking(
pii_entities_config={},
presidio_language="es",
mock_testing=True
pii_entities_config={}, presidio_language="es", mock_testing=True
)
test_text_es = "Mi número de teléfono es +34 912 345 678"
analyze_request_es = presidio_guardrail_es._get_presidio_analyze_request_payload(
text=test_text_es,
presidio_config=None,
request_data={}
text=test_text_es, presidio_config=None, request_data={}
)
# Verify the language is set to Spanish
assert analyze_request_es["language"] == "es"
assert analyze_request_es["text"] == test_text_es
# Test default language (English) when not specified
presidio_guardrail_default = _OPTIONAL_PresidioPIIMasking(
pii_entities_config={},
mock_testing=True
pii_entities_config={}, mock_testing=True
)
test_text_en = "My phone number is +1 555-123-4567"
analyze_request_default = presidio_guardrail_default._get_presidio_analyze_request_payload(
text=test_text_en,
presidio_config=None,
request_data={}
analyze_request_default = (
presidio_guardrail_default._get_presidio_analyze_request_payload(
text=test_text_en, presidio_config=None, request_data={}
)
)
# Verify the language defaults to English
assert analyze_request_default["language"] == "en"
assert analyze_request_default["text"] == test_text_en
@ -562,36 +560,30 @@ async def test_presidio_language_configuration():
async def test_presidio_language_configuration_with_per_request_override():
"""Test that per-request language configuration overrides the default configured language"""
litellm._turn_on_debug()
# Set up guardrail with German as default language
presidio_guardrail = _OPTIONAL_PresidioPIIMasking(
pii_entities_config={},
presidio_language="de",
mock_testing=True
pii_entities_config={}, presidio_language="de", mock_testing=True
)
test_text = "Test text with PII"
# Test with per-request config overriding the default language
presidio_config = PresidioPerRequestConfig(language="fr")
analyze_request = presidio_guardrail._get_presidio_analyze_request_payload(
text=test_text,
presidio_config=presidio_config,
request_data={}
text=test_text, presidio_config=presidio_config, request_data={}
)
# Verify the per-request language (French) overrides the default (German)
assert analyze_request["language"] == "fr"
assert analyze_request["text"] == test_text
# Test without per-request config - should use default language
analyze_request_default = presidio_guardrail._get_presidio_analyze_request_payload(
text=test_text,
presidio_config=None,
request_data={}
text=test_text, presidio_config=None, request_data={}
)
# Verify the default language (German) is used
assert analyze_request_default["language"] == "de"
assert analyze_request_default["text"] == test_text

622
tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py
View file

@ -1608,3 +1608,625 @@ async def test_anonymize_text_http_error_status():
output_parse_pii=False,
masked_entity_count={},
)
@pytest.mark.asyncio
async def test_pii_tokens_stored_in_metadata_not_top_level(presidio_guardrail):
"""
Regression test: pii_tokens must be stored in data['metadata']['pii_tokens'],
NOT in data['pii_tokens']. Storing at the top level leaks the field to LLM
providers like Anthropic, which reject unknown fields with
'pii_tokens: Extra inputs are not permitted'.
"""
guardrail = _OPTIONAL_PresidioPIIMasking(
mock_testing=True,
output_parse_pii=True,
pii_entities_config={
PiiEntityType.PERSON: PiiAction.MASK,
PiiEntityType.PHONE_NUMBER: PiiAction.MASK,
},
)
mock_user_api_key = UserAPIKeyAuth(api_key="test-key")
mock_cache = DualCache()
test_data = {
"messages": [
{"role": "user", "content": "My name is John and my phone is 555-123-4567"}
],
"model": "claude-haiku-4-5-20251001",
"metadata": {},
}
async def mock_check_pii(text, output_parse_pii, presidio_config, request_data):
# Simulate PII masking with token storage (mimics real anonymize_text behavior)
if request_data is not None and output_parse_pii:
if "metadata" not in request_data:
request_data["metadata"] = {}
if "pii_tokens" not in request_data["metadata"]:
request_data["metadata"]["pii_tokens"] = {}
pii_tokens = request_data["metadata"]["pii_tokens"]
seq = len(pii_tokens) + 1
token = f"<PERSON_{seq}>"
pii_tokens[token] = "John"
text = text.replace("John", token)
return text
guardrail.check_pii = mock_check_pii
result = await guardrail.async_pre_call_hook(
user_api_key_dict=mock_user_api_key,
cache=mock_cache,
data=test_data,
call_type="completion",
)
# pii_tokens must NOT be at the top level of data (would leak to providers)
assert "pii_tokens" not in result, (
"pii_tokens must not be a top-level key in request data — "
"it would leak to LLM providers and cause 'Extra inputs are not permitted' errors"
)
# pii_tokens must be inside metadata (safe from provider leakage)
assert "metadata" in result
assert "pii_tokens" in result["metadata"]
assert len(result["metadata"]["pii_tokens"]) > 0
@pytest.mark.asyncio
async def test_pii_tokens_in_metadata_used_for_unmasking():
"""
Regression test: _process_response_for_pii must read pii_tokens from
data['metadata']['pii_tokens'] and correctly unmask the response.
"""
guardrail = _OPTIONAL_PresidioPIIMasking(
mock_testing=True,
output_parse_pii=True,
)
token_key = "<PERSON_1>"
request_data = {
"model": "claude-haiku-4-5-20251001",
"metadata": {"pii_tokens": {token_key: "John"}},
}
response = ModelResponse(
choices=[
Choices(
message=Message(
role="assistant",
content=f"Hello {token_key}, how can I help you?",
),
index=0,
finish_reason="stop",
)
]
)
await guardrail._process_response_for_pii(
response=response,
request_data=request_data,
mode="unmask",
)
assert response.choices[0].message.content == "Hello John, how can I help you?"
@pytest.mark.parametrize(
"initial_hook",
["pre_call", "during_call", "pre_mcp_call"],
)
def test_event_hook_auto_expansion_for_all_string_hooks(initial_hook):
"""
Regression test: when output_parse_pii is True, the guardrail must add
'post_call' to event_hook regardless of the initial string hook value,
not just when it's 'pre_call'.
"""
guardrail = _OPTIONAL_PresidioPIIMasking(
mock_testing=True,
output_parse_pii=True,
event_hook=initial_hook,
)
assert isinstance(guardrail.event_hook, list)
assert initial_hook in guardrail.event_hook
assert "post_call" in guardrail.event_hook
def test_event_hook_no_expansion_when_already_post_call():
"""post_call alone should stay as-is — no expansion needed."""
guardrail = _OPTIONAL_PresidioPIIMasking(
mock_testing=True,
output_parse_pii=True,
event_hook="post_call",
)
# Should remain a string "post_call", not expanded to a list
assert guardrail.event_hook == "post_call"
@pytest.mark.asyncio
async def test_metadata_none_does_not_crash():
"""
Regression test: if metadata is explicitly None in request_data,
the guardrail must not crash with TypeError on the write or read path.
"""
guardrail = _OPTIONAL_PresidioPIIMasking(
mock_testing=True,
output_parse_pii=True,
)
token_key = "<PERSON_1>"
# metadata explicitly None — must not crash
request_data = {
"model": "gpt-3.5-turbo",
"metadata": None,
}
response = ModelResponse(
choices=[
Choices(
message=Message(
role="assistant",
content=f"Hello {token_key}, how can I help you?",
),
index=0,
finish_reason="stop",
)
]
)
# Should not raise TypeError
await guardrail._process_response_for_pii(
response=response,
request_data=request_data,
mode="unmask",
)
# 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?"
)
# ---------------------------------------------------------------------------
# Tests for sequential-numbered token unmasking in _unmask_pii_text
# ---------------------------------------------------------------------------
def test_unmask_exact_match_with_sequential_tokens():
"""
Normal unmasking: LLM echoes numbered tokens verbatim → original PII restored.
"""
from litellm.proxy.guardrails.guardrail_hooks.presidio import (
_OPTIONAL_PresidioPIIMasking,
)
pii_tokens = {
"<PERSON_1>": "John Smith",
"<PHONE_NUMBER_1>": "555-123-4567",
}
text = "Hello <PERSON_1>, your number is <PHONE_NUMBER_1>."
result = _OPTIONAL_PresidioPIIMasking._unmask_pii_text(text, pii_tokens)
assert result == "Hello John Smith, your number is 555-123-4567."
def test_unmask_multiple_same_entity_type():
"""
Two phone numbers get distinct numbered tokens and unmask correctly.
"""
from litellm.proxy.guardrails.guardrail_hooks.presidio import (
_OPTIONAL_PresidioPIIMasking,
)
pii_tokens = {
"<PHONE_NUMBER_1>": "555-111-0000",
"<PHONE_NUMBER_2>": "555-222-0000",
}
text = "Call <PHONE_NUMBER_1> or <PHONE_NUMBER_2>."
result = _OPTIONAL_PresidioPIIMasking._unmask_pii_text(text, pii_tokens)
assert result == "Call 555-111-0000 or 555-222-0000."
def test_unmask_graceful_degradation():
"""
If the LLM doesn't echo the token back, the numbered label stays
in the output — clean and readable, not garbage hex.
"""
from litellm.proxy.guardrails.guardrail_hooks.presidio import (
_OPTIONAL_PresidioPIIMasking,
)
pii_tokens = {
"<PERSON_1>": "John",
}
# LLM paraphrased instead of echoing the token
text = "I see you provided a name."
result = _OPTIONAL_PresidioPIIMasking._unmask_pii_text(text, pii_tokens)
# No change — no garbage, just clean text
assert result == text
# ---------------------------------------------------------------------------
# Fix 1: Position bug — reverse sort + original text coordinates
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_anonymize_text_multiple_items_position_correctness():
"""
Regression test: when multiple PII items exist, coordinates reference the
ORIGINAL text. Processing in reverse order prevents coordinate drift.
"""
guardrail = _OPTIONAL_PresidioPIIMasking(
presidio_analyzer_api_base="http://test-analyzer/",
presidio_anonymizer_api_base="http://test-anonymizer/",
mock_testing=False,
)
# "Call John at 555-123-4567"
# "John" at [5:9], "555-123-4567" at [13:25]
anonymizer_response = {
"text": "Call <PERSON> at <PHONE_NUMBER>",
"items": [
{
"start": 5,
"end": 9,
"entity_type": "PERSON",
"text": "<PERSON>",
"operator": "replace",
},
{
"start": 13,
"end": 25,
"entity_type": "PHONE_NUMBER",
"text": "<PHONE_NUMBER>",
"operator": "replace",
},
],
}
mock_iterator = _make_mock_session_iterator(anonymizer_response)
request_data = {"metadata": {}}
with patch.object(guardrail, "_get_session_iterator", mock_iterator):
result = await guardrail.anonymize_text(
text="Call John at 555-123-4567",
analyze_results=[
{"start": 5, "end": 9, "entity_type": "PERSON", "score": 0.9},
{"start": 13, "end": 25, "entity_type": "PHONE_NUMBER", "score": 0.95},
],
output_parse_pii=True,
masked_entity_count={},
request_data=request_data,
)
pii_tokens = request_data["metadata"]["pii_tokens"]
# Verify tokens captured the correct ORIGINAL text values
person_token = [k for k in pii_tokens if "PERSON" in k][0]
phone_token = [k for k in pii_tokens if "PHONE" in k][0]
assert pii_tokens[person_token] == "John"
assert pii_tokens[phone_token] == "555-123-4567"
# Verify both PII values are masked in the result
assert "John" not in result
assert "555-123-4567" not in result
# ---------------------------------------------------------------------------
# Fix 2: Anthropic native dict response handling
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_anthropic_native_response_unmasking():
"""
Anthropic native dict responses (type='message') should be unmasked
when output_parse_pii is enabled.
"""
guardrail = _OPTIONAL_PresidioPIIMasking(
mock_testing=True,
output_parse_pii=True,
)
request_data = {
"model": "claude-3-haiku",
"metadata": {
"pii_tokens": {
"<PERSON_1>": "John Smith",
"<PHONE_NUMBER_1>": "555-123-4567",
}
},
}
anthropic_response = {
"type": "message",
"id": "msg_123",
"model": "claude-3-haiku",
"role": "assistant",
"content": [
{
"type": "text",
"text": "Hello <PERSON_1>, your number is <PHONE_NUMBER_1>.",
}
],
"stop_reason": "end_turn",
"usage": {"input_tokens": 10, "output_tokens": 20},
}
mock_user_api_key = UserAPIKeyAuth(api_key="test-key")
result = await guardrail.async_post_call_success_hook(
data=request_data,
user_api_key_dict=mock_user_api_key,
response=anthropic_response,
)
assert result["content"][0]["text"] == (
"Hello John Smith, your number is 555-123-4567."
)
@pytest.mark.asyncio
async def test_anthropic_native_response_masking():
"""
Anthropic native dict responses should be masked when
apply_to_output is enabled.
"""
guardrail = _OPTIONAL_PresidioPIIMasking(
mock_testing=True,
apply_to_output=True,
)
async def mock_check_pii(text, output_parse_pii, presidio_config, request_data):
return text.replace("John Smith", "[PERSON]").replace("555-123-4567", "[PHONE]")
guardrail.check_pii = mock_check_pii
anthropic_response = {
"type": "message",
"id": "msg_123",
"model": "claude-3-haiku",
"role": "assistant",
"content": [{"type": "text", "text": "Hello John Smith, call 555-123-4567."}],
"stop_reason": "end_turn",
"usage": {"input_tokens": 10, "output_tokens": 20},
}
mock_user_api_key = UserAPIKeyAuth(api_key="test-key")
result = await guardrail.async_post_call_success_hook(
data={},
user_api_key_dict=mock_user_api_key,
response=anthropic_response,
)
assert "[PERSON]" in result["content"][0]["text"]
assert "[PHONE]" in result["content"][0]["text"]
assert "John Smith" not in result["content"][0]["text"]
@pytest.mark.asyncio
async def test_anthropic_native_response_non_text_blocks_untouched():
"""
Non-text blocks (tool_use, thinking) in Anthropic responses
should be left untouched during unmasking.
"""
guardrail = _OPTIONAL_PresidioPIIMasking(
mock_testing=True,
output_parse_pii=True,
)
request_data = {
"model": "claude-3-haiku",
"metadata": {"pii_tokens": {"<PERSON_1>": "John"}},
}
anthropic_response = {
"type": "message",
"id": "msg_123",
"content": [
{"type": "text", "text": "Hello <PERSON_1>"},
{
"type": "tool_use",
"id": "call_1",
"name": "search",
"input": {"q": "test"},
},
],
"role": "assistant",
"stop_reason": "end_turn",
"usage": {"input_tokens": 10, "output_tokens": 20},
}
mock_user_api_key = UserAPIKeyAuth(api_key="test-key")
result = await guardrail.async_post_call_success_hook(
data=request_data,
user_api_key_dict=mock_user_api_key,
response=anthropic_response,
)
assert result["content"][0]["text"] == "Hello John"
assert result["content"][1]["type"] == "tool_use"
assert result["content"][1]["name"] == "search"
# ---------------------------------------------------------------------------
# Fix 3: Anthropic native SSE streaming — bytes passthrough
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_streaming_bytes_chunks_are_yielded_not_discarded():
"""
Regression test: bytes chunks (Anthropic native SSE) should be yielded
through the streaming hook, not silently discarded.
"""
guardrail = _OPTIONAL_PresidioPIIMasking(
mock_testing=True,
apply_to_output=True,
)
byte_chunk = b'data: {"type":"content_block_delta","delta":{"text":"Hello"}}\n\n'
async def mock_stream():
yield byte_chunk
mock_user_api_key = UserAPIKeyAuth(api_key="test-key")
chunks = []
async for chunk in guardrail.async_post_call_streaming_iterator_hook(
user_api_key_dict=mock_user_api_key,
response=mock_stream(),
request_data={},
):
chunks.append(chunk)
assert any(
isinstance(c, bytes) for c in chunks
), "bytes chunks must not be discarded"
assert byte_chunk in chunks
@pytest.mark.asyncio
async def test_streaming_unmask_path_bytes_passthrough():
"""
Bytes chunks in the unmasking path should also pass through.
"""
guardrail = _OPTIONAL_PresidioPIIMasking(
mock_testing=True,
output_parse_pii=True,
)
byte_chunk = b'data: {"type":"content_block_delta"}\n\n'
request_data = {
"metadata": {"pii_tokens": {"<PERSON_1>": "John"}},
}
async def mock_stream():
yield byte_chunk
mock_user_api_key = UserAPIKeyAuth(api_key="test-key")
chunks = []
async for chunk in guardrail.async_post_call_streaming_iterator_hook(
user_api_key_dict=mock_user_api_key,
response=mock_stream(),
request_data=request_data,
):
chunks.append(chunk)
assert len(chunks) == 1
assert chunks[0] == byte_chunk
# ---------------------------------------------------------------------------
# Fix 4: apply_guardrail unmask path for input_type="response"
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_apply_guardrail_unmask_on_response():
"""
When input_type is 'response' and pii_tokens exist, apply_guardrail
should unmask text instead of masking it.
"""
guardrail = _OPTIONAL_PresidioPIIMasking(
guardrail_name="test_presidio",
output_parse_pii=True,
mock_testing=True,
)
request_data = {
"model": "gpt-4o",
"metadata": {
"pii_tokens": {
"<PERSON_1>": "John Smith",
"<PHONE_NUMBER_1>": "555-123-4567",
}
},
}
inputs = {
"texts": [
"Hello <PERSON_1>, your number is <PHONE_NUMBER_1>.",
]
}
result = await guardrail.apply_guardrail(
inputs=inputs,
request_data=request_data,
input_type="response",
)
assert result["texts"][0] == "Hello John Smith, your number is 555-123-4567."
@pytest.mark.asyncio
async def test_apply_guardrail_masks_on_request():
"""
When input_type is 'request', apply_guardrail should mask as before.
"""
guardrail = _OPTIONAL_PresidioPIIMasking(
guardrail_name="test_presidio",
output_parse_pii=True,
mock_testing=True,
)
async def mock_check_pii(text, output_parse_pii, presidio_config, request_data):
return text.replace("John Smith", "<PERSON>")
guardrail.check_pii = mock_check_pii
result = await guardrail.apply_guardrail(
inputs={"texts": ["Hello John Smith"]},
request_data={"model": "gpt-4o", "metadata": {}},
input_type="request",
)
assert "<PERSON>" in result["texts"][0]
assert "John Smith" not in result["texts"][0]
@pytest.mark.asyncio
async def test_apply_to_output_streaming_bytes_only_logs_warning():
"""
Regression test: when apply_to_output=True and the stream contains only
bytes chunks (Anthropic native SSE), output masking is skipped.
A warning must be logged so operators are aware.
"""
guardrail = _OPTIONAL_PresidioPIIMasking(
mock_testing=True,
apply_to_output=True,
)
byte_chunks = [
b'data: {"type":"content_block_delta","delta":{"text":"Hello"}}\n\n',
b'data: {"type":"content_block_delta","delta":{"text":" world"}}\n\n',
]
async def mock_stream():
for b in byte_chunks:
yield b
mock_user_api_key = UserAPIKeyAuth(api_key="test-key")
collected = []
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(),
request_data={},
):
collected.append(chunk)
# All bytes should be yielded through
assert len(collected) == len(byte_chunks)
for original, received in zip(byte_chunks, collected):
assert original == received
# Warning must be logged about skipped masking
mock_logger.warning.assert_called_once()
warning_msg = mock_logger.warning.call_args[0][0]
assert "Output PII masking was skipped" in warning_msg