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

fix presidio output_parse_pii token handoff

Browse source 
Keep pii token mappings request-scoped across guardrail pre/post paths so output_parse_pii can reliably unmask responses. Add regression tests for litellm_metadata token fallback and response behavior when tokens are missing.

Made-with: Cursor
This commit is contained in:
Milan 2026-04-03 13:15:05 +02:00
parent d4a3a5e530
commit 672f939631
No known key found for this signature in database
2 changed files with 135 additions and 21 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -504,13 +504,9 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
"This may indicate a missing caller update."
)
request_data = {}
# 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"]
pii_tokens = self._get_or_create_request_scoped_pii_tokens(
request_data
)
# Append a sequential number to make each token unique
# per request, so unmasking maps back to the correct
@ -935,6 +931,56 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
)
return response
def _get_or_create_request_scoped_pii_tokens(
self, request_data: Dict[str, Any]
) -> Dict[str, str]:
"""
Return request-scoped token map and ensure it is anchored in all
metadata locations used in pre/post guardrail hooks.
"""
metadata = request_data.get("metadata")
if not isinstance(metadata, dict):
metadata = {}
request_data["metadata"] = metadata
pii_tokens = metadata.get("pii_tokens")
if not isinstance(pii_tokens, dict):
pii_tokens = {}
metadata["pii_tokens"] = pii_tokens
litellm_metadata = request_data.get("litellm_metadata")
if isinstance(litellm_metadata, dict):
litellm_metadata["pii_tokens"] = pii_tokens
return cast(Dict[str, str], pii_tokens)
def _get_request_scoped_pii_tokens(
self, request_data: Optional[Dict[str, Any]]
) -> Dict[str, str]:
"""
Read request-scoped pii token map with fallback across known metadata paths.
"""
if not request_data:
return {}
metadata = request_data.get("metadata")
if isinstance(metadata, dict):
pii_tokens = metadata.get("pii_tokens")
if isinstance(pii_tokens, dict):
litellm_metadata = request_data.get("litellm_metadata")
if isinstance(litellm_metadata, dict):
litellm_metadata["pii_tokens"] = pii_tokens
return cast(Dict[str, str], pii_tokens)
litellm_metadata = request_data.get("litellm_metadata")
if isinstance(litellm_metadata, dict):
pii_tokens = litellm_metadata.get("pii_tokens")
if isinstance(pii_tokens, dict):
return cast(Dict[str, str], pii_tokens)
return {}
@staticmethod
def _unmask_pii_text(text: str, pii_tokens: Dict[str, str]) -> str:
"""
@ -980,11 +1026,10 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
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", {})
pii_tokens = self._get_request_scoped_pii_tokens(request_data)
if not pii_tokens and mode == "unmask":
verbose_proxy_logger.debug(
"No pii_tokens in metadata for Anthropic response unmask"
"No pii_tokens found for Anthropic response unmask"
)
presidio_config = self.get_presidio_settings_from_request_data(
request_data or {}
@ -1022,11 +1067,10 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
Helper to recursively process a ModelResponse for PII.
Handles all choices and tool calls.
"""
metadata = (request_data.get("metadata") or {}) if request_data else {}
pii_tokens = metadata.get("pii_tokens", {})
pii_tokens = self._get_request_scoped_pii_tokens(request_data)
if not pii_tokens and mode == "unmask":
verbose_proxy_logger.debug(
"No pii_tokens found in request_data['metadata'] — nothing to unmask"
"No pii_tokens found in request_data — nothing to unmask"
)
presidio_config = self.get_presidio_settings_from_request_data(
request_data or {}
@ -1253,11 +1297,10 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
yield chunk
return
metadata = (request_data.get("metadata") or {}) if request_data else {}
pii_tokens = metadata.get("pii_tokens", {})
pii_tokens = self._get_request_scoped_pii_tokens(request_data)
if not pii_tokens and request_data:
verbose_proxy_logger.debug(
"No pii_tokens in request_data['metadata'] for streaming unmask path"
"No pii_tokens found in request_data for streaming unmask path"
)
if not (self.output_parse_pii and pii_tokens):
async for chunk in response:
@ -1317,13 +1360,28 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
# 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", {})
pii_tokens = self._get_request_scoped_pii_tokens(request_data)
new_texts = []
if input_type == "response" and pii_tokens:
for text in texts:
new_texts.append(self._unmask_pii_text(text, pii_tokens))
if input_type == "response":
if pii_tokens:
for text in texts:
new_texts.append(self._unmask_pii_text(text, pii_tokens))
elif self.output_parse_pii:
verbose_proxy_logger.debug(
"output_parse_pii enabled but no pii_tokens found on response path; "
"leaving response text unchanged"
)
new_texts = list(texts)
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)
else:
for text in texts:
modified_text = await self.check_pii(

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

@ -2161,6 +2161,62 @@ async def test_apply_guardrail_unmask_on_response():
assert result["texts"][0] == "Hello John Smith, your number is 555-123-4567."
@pytest.mark.asyncio
async def test_apply_guardrail_unmask_reads_tokens_from_litellm_metadata():
"""
Regression test: if pii_tokens are present on litellm_metadata (but not metadata),
response unmasking should still work.
"""
guardrail = _OPTIONAL_PresidioPIIMasking(
guardrail_name="test_presidio",
output_parse_pii=True,
mock_testing=True,
)
request_data = {
"model": "gpt-4o",
"litellm_metadata": {
"pii_tokens": {
"<PERSON_1>": "Jane Doe",
}
},
}
result = await guardrail.apply_guardrail(
inputs={"texts": ["Hello <PERSON_1>"]},
request_data=request_data,
input_type="response",
)
assert result["texts"][0] == "Hello Jane Doe"
@pytest.mark.asyncio
async def test_apply_guardrail_response_no_tokens_does_not_remask():
"""
Regression test: for output_parse_pii response path with no tokens,
text should be returned unchanged (no re-masking pass).
"""
guardrail = _OPTIONAL_PresidioPIIMasking(
guardrail_name="test_presidio",
output_parse_pii=True,
mock_testing=True,
)
async def fail_if_called(*args, **kwargs):
raise AssertionError("check_pii should not run on response path without tokens")
guardrail.check_pii = fail_if_called
result = await guardrail.apply_guardrail(
inputs={"texts": ["Nice to meet you, <PERSON_1>!"]},
request_data={"model": "gpt-4o", "metadata": {}},
input_type="response",
)
assert result["texts"][0] == "Nice to meet you, <PERSON_1>!"
@pytest.mark.asyncio
async def test_apply_guardrail_masks_on_request():
"""