scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-24 00:52:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 90

fix(presidio): restore PII unmasking in streaming responses

Browse source 
When output_parse_pii is enabled with presidio_filter_scope=both, a separate
apply_to_output callback was re-masking the stream after unmask ran. Skip
registering that callback when output_parse_pii is on, mirror pii_tokens across
metadata containers, and passthrough apply_to_output streaming when tokens exist.

Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
shivam 2026-05-19 19:09:29 -07:00
parent e59e34bed3
commit 9d7fc90c54
No known key found for this signature in database
3 changed files with 190 additions and 17 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -482,6 +482,40 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
)
return redacted_text["text"]
def _ensure_pii_tokens(self, request_data: dict) -> Dict[str, str]:
"""
Return the shared pii_tokens map for this request, stored on both
metadata and litellm_metadata so post-call unmasking works regardless
of which container downstream hooks read.
"""
tokens: Optional[Dict[str, str]] = None
for meta_key in ("metadata", "litellm_metadata"):
meta = request_data.get(meta_key)
if isinstance(meta, dict) and isinstance(meta.get("pii_tokens"), dict):
tokens = meta["pii_tokens"]
break
if tokens is None:
tokens = {}
for meta_key in ("metadata", "litellm_metadata"):
if not isinstance(request_data.get(meta_key), dict):
request_data[meta_key] = {}
request_data[meta_key]["pii_tokens"] = tokens
return tokens
@staticmethod
def _get_pii_tokens_from_request_data(
request_data: Optional[Dict],
) -> Dict[str, str]:
if not request_data:
return {}
for meta_key in ("metadata", "litellm_metadata"):
meta = request_data.get(meta_key)
if isinstance(meta, dict):
tokens = meta.get("pii_tokens")
if isinstance(tokens, dict) and tokens:
return tokens
return {}
def _finalize_presidio_anonymize_numbered_tokens(
self,
text: str,
@ -501,11 +535,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
"This may indicate a missing caller update."
)
request_data = {}
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._ensure_pii_tokens(request_data)
# Assign sequence numbers in forward (left-to-right) order so
# that <PERSON_1> is the first entity in the text, etc.
@ -1001,8 +1031,7 @@ 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_pii_tokens_from_request_data(request_data)
if not pii_tokens and mode == "unmask":
verbose_proxy_logger.debug(
"No pii_tokens in metadata for Anthropic response unmask"
@ -1043,11 +1072,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_pii_tokens_from_request_data(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 metadata — nothing to unmask"
)
presidio_config = self.get_presidio_settings_from_request_data(
request_data or {}
@ -1292,17 +1320,21 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
The base class declares ModelResponseStream only.
"""
if self.apply_to_output:
pii_tokens = self._get_pii_tokens_from_request_data(request_data)
if pii_tokens:
async for chunk in response:
yield chunk
return
async for chunk in self._stream_apply_output_masking(
response, request_data
):
yield chunk
return
metadata = (request_data.get("metadata") or {}) if request_data else {}
pii_tokens = metadata.get("pii_tokens", {})
pii_tokens = self._get_pii_tokens_from_request_data(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 in request_data for streaming unmask path"
)
if not (self.output_parse_pii and pii_tokens):
async for chunk in response:
@ -1362,8 +1394,7 @@ 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_pii_tokens_from_request_data(request_data)
new_texts = []
if input_type == "response" and pii_tokens:
@ -1394,3 +1425,17 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
self.presidio_entities_deny_list = (
litellm_params.presidio_entities_deny_list
)
if litellm_params.output_parse_pii is not None:
self.output_parse_pii = litellm_params.output_parse_pii or False
if (self.output_parse_pii or self.apply_to_output) and not getattr(
self, "logging_only", False
):
current_hook = self.event_hook
if isinstance(current_hook, str) and current_hook != "post_call":
self.event_hook = cast(
List[GuardrailEventHooks], [current_hook, "post_call"]
)
elif isinstance(current_hook, list) and "post_call" not in current_hook:
self.event_hook = cast(
List[GuardrailEventHooks], current_hook + ["post_call"]
)

2
litellm/proxy/guardrails/guardrail_initializers.py
View file

@ -112,7 +112,7 @@ def initialize_presidio(litellm_params: LitellmParams, guardrail: Guardrail):
event_hook=GuardrailEventHooks.post_call.value,
)
if run_output:
if run_output and not litellm_params.output_parse_pii:
output_callback = _make_presidio_callback(
apply_to_output=True,
event_hook=GuardrailEventHooks.post_call.value,

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

@ -21,7 +21,14 @@ from litellm.proxy.guardrails.guardrail_hooks.presidio import (
)
from litellm.exceptions import GuardrailRaisedException
from litellm.types.guardrails import LitellmParams, PiiAction, PiiEntityType
from litellm.types.utils import Choices, Message, ModelResponse
from litellm.types.utils import (
Choices,
Delta,
Message,
ModelResponse,
ModelResponseStream,
StreamingChoices,
)
def _make_mock_session_iterator(
@ -797,6 +804,61 @@ async def test_presidio_filter_scope_initializer(monkeypatch):
assert any(not c.apply_to_output for c in created)
assert any(c.apply_to_output for c in created)
# both + output_parse_pii -> unmask only; no apply_to_output callback
created.clear()
params_both_unmask = LitellmParams(
guardrail="presidio",
mode="pre_call",
presidio_filter_scope="both",
output_parse_pii=True,
)
cb = initialize_presidio(params_both_unmask, guardrail_dict)
assert len(created) == 2
assert all(not c.apply_to_output for c in created)
@pytest.mark.asyncio
async def test_apply_to_output_streaming_skips_remask_when_pii_tokens_present():
"""
When input was masked with numbered tokens, the apply_to_output streaming
hook must not re-mask the already-unmasked response.
"""
guardrail = _OPTIONAL_PresidioPIIMasking(
guardrail_name="test_presidio",
apply_to_output=True,
event_hook="post_call",
mock_testing=True,
)
request_data = {
"metadata": {
"pii_tokens": {"<EMAIL_ADDRESS_1>": "shivam@uni.minerva.edu"},
}
}
async def _fake_stream():
yield ModelResponseStream(
choices=[
StreamingChoices(
delta=Delta(content="Your email is shivam@uni.minerva.edu"),
finish_reason="stop",
index=0,
)
]
)
chunks = []
async for chunk in guardrail.async_post_call_streaming_iterator_hook(
user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"),
response=_fake_stream(),
request_data=request_data,
):
chunks.append(chunk)
assert len(chunks) == 1
assert "shivam@uni.minerva.edu" in chunks[0].choices[0].delta.content
assert "<EMAIL_ADDRESS" not in chunks[0].choices[0].delta.content
@pytest.mark.asyncio
async def test_empty_content_handling(
@ -2258,6 +2320,72 @@ 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_proxy_post_call_unmask_when_litellm_metadata_shadows_guardrails():
"""
Regression: guardrails live in metadata but a non-empty litellm_metadata
without guardrails must not prevent Presidio post-call unmasking.
"""
import litellm
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.utils import ProxyLogging
litellm.callbacks = []
post_guard = _OPTIONAL_PresidioPIIMasking(
mock_testing=True,
guardrail_name="presidio-g",
output_parse_pii=True,
event_hook="post_call",
)
litellm.logging_callback_manager.add_litellm_callback(post_guard)
data = {
"model": "gpt-4",
"metadata": {
"guardrails": ["presidio-g"],
"pii_tokens": {"<PERSON_1>": "John Smith"},
},
"litellm_metadata": {"user_api_key_user_id": "user-1"},
}
response = ModelResponse(
choices=[
Choices(
message=Message(
role="assistant",
content="Hello <PERSON_1>, how can I help?",
),
index=0,
finish_reason="stop",
)
]
)
proxy_logging = ProxyLogging(user_api_key_cache=None)
result = await proxy_logging.post_call_success_hook(
data=data,
response=response,
user_api_key_dict=UserAPIKeyAuth(
api_key="test-key", request_route="/chat/completions"
),
)
assert result.choices[0].message.content == "Hello John Smith, how can I help?"
def test_pii_tokens_mirrored_to_litellm_metadata():
"""pii_tokens must be readable from either metadata container after masking."""
guardrail = _OPTIONAL_PresidioPIIMasking(mock_testing=True)
request_data: dict = {"litellm_metadata": {"user_api_key_user_id": "u1"}}
tokens = guardrail._ensure_pii_tokens(request_data)
tokens["<PERSON_1>"] = "John Smith"
assert request_data["metadata"]["pii_tokens"]["<PERSON_1>"] == "John Smith"
assert request_data["litellm_metadata"]["pii_tokens"]["<PERSON_1>"] == "John Smith"
assert guardrail._get_pii_tokens_from_request_data(request_data)["<PERSON_1>"] == (
"John Smith"
)
@pytest.mark.asyncio
async def test_apply_guardrail_masks_on_request():
"""