diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index 1e4621622b3..5f2b6ef3e20 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -991,6 +991,23 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): return filtered_results + def _exclude_restored_request_values( + self, + text: str, + analyze_results: list[PresidioAnalyzeResponseItem], + request_data: dict, + ) -> list[PresidioAnalyzeResponseItem]: + restored_values: Final = set(((request_data.get("metadata") or {}).get("pii_tokens") or {}).values()) + if not restored_values: + return analyze_results + + def _span_is_restored(item: PresidioAnalyzeResponseItem) -> bool: + start: Final = item.get("start") + end: Final = item.get("end") + return start is not None and end is not None and text[start:end] in restored_values + + return [item for item in analyze_results if not _span_is_restored(item)] + def raise_exception_if_blocked_entities_detected( self, analyze_results: list[PresidioAnalyzeResponseItem] | _PresidioAnonymizeResponse ) -> None: @@ -1052,6 +1069,13 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): #################################################### self.raise_exception_if_blocked_entities_detected(analyze_results=analyze_results) + if output_parse_pii is False and isinstance(analyze_results, list): + analyze_results = self._exclude_restored_request_values( + text=text, + analyze_results=analyze_results, + request_data=request_data, + ) + # Then anonymize the text using the analysis results anonymized_text: Final = await self.anonymize_text( text=text, diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_presidio.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_presidio.py index cc0d2a8008f..c8b79d87e02 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_presidio.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_presidio.py @@ -5480,3 +5480,165 @@ def test_get_session_iterator_reuses_one_session_per_background_loop( assert len(sessions) == 10 assert all(session is sessions[0] for session in sessions) assert presidio_guardrail._http_session is None + + +def _presidio_fake_server_detections(text: str) -> list[dict]: + detections: Final = [] + for match in re.finditer(r"[\w.+-]+@[\w-]+(?:\.[\w-]+)+", text): + detections.append( + { + "entity_type": "EMAIL_ADDRESS", + "start": match.start(), + "end": match.end(), + "score": 0.9, + } + ) + for match in re.finditer(r"GB\d{2}[A-Z]{4}\d{14}", text): + detections.append( + { + "entity_type": "IBAN_CODE", + "start": match.start(), + "end": match.end(), + "score": 0.9, + } + ) + return detections + + +def _presidio_fake_server_iterator(): + """Presidio stand-in that masks only the analyzer_results it is handed.""" + + @asynccontextmanager + async def mock_iterator(): + class MockResponse: + def __init__(self, payload): + self.payload = payload + self.status = 200 + self.content_type = "application/json" + self.headers = {"Content-Type": "application/json"} + + async def text(self): + return json.dumps(self.payload) + + async def json(self): + return self.payload + + async def __aenter__(self): + return self + + async def __aexit__(self, *args): + pass + + class MockSession: + closed = False + + def post(self, url, json=None, **kwargs): + if url.endswith("analyze"): + return MockResponse(_presidio_fake_server_detections(json["text"])) + text = json["text"] + for span in sorted(json.get("analyzer_results") or [], key=lambda i: i["start"], reverse=True): + text = text[: span["start"]] + f"<{span['entity_type']}>" + text[span["end"] :] + return MockResponse({"text": text, "items": json.get("analyzer_results") or []}) + + async def __aenter__(self): + return self + + async def __aexit__(self, *args): + pass + + yield MockSession() + + return mock_iterator + + +@pytest.mark.asyncio +async def test_output_masking_does_not_remask_restored_values(): + """ + A response that echoes numbered input tokens must come back restored, + even when an output-masking callback runs after the restore step, while + PII the model generated itself must still be masked. + """ + from litellm.types.utils import ModelResponseStream + + restore_guardrail = OPTIONAL_PresidioPIIMasking( + output_parse_pii=True, + presidio_analyzer_api_base="http://analyzer:3000/", + presidio_anonymizer_api_base="http://anonymizer:3000/", + pii_entities_config={ + PiiEntityType.EMAIL_ADDRESS: PiiAction.MASK, + "IBAN_CODE": PiiAction.MASK, + }, + _callback_role="restore", + ) + mask_guardrail = OPTIONAL_PresidioPIIMasking( + apply_to_output=True, + presidio_analyzer_api_base="http://analyzer:3000/", + presidio_anonymizer_api_base="http://anonymizer:3000/", + pii_entities_config={ + PiiEntityType.EMAIL_ADDRESS: PiiAction.MASK, + "IBAN_CODE": PiiAction.MASK, + }, + ) + user_api_key = UserAPIKeyAuth(api_key="test-key") + request_data: Final = { + "messages": [ + { + "role": "user", + "content": "Repeat this line verbatim: Contact , account .", + } + ], + "model": "fake-model", + "metadata": { + "pii_tokens": { + "": "test.person@contoso-example.com", + "": "GB82WEST12345698765432", + } + }, + } + echo = "Contact , account . New contact: other.person@example.org." + expected = "Contact test.person@contoso-example.com, account GB82WEST12345698765432. New contact: ." + + with ( + patch.object(restore_guardrail, "_get_session_iterator", _presidio_fake_server_iterator()), + patch.object(mask_guardrail, "_get_session_iterator", _presidio_fake_server_iterator()), + ): + response = ModelResponse( + choices=[Choices(message=Message(role="assistant", content=echo), index=0, finish_reason="stop")] + ) + await restore_guardrail.async_post_call_success_hook( + data=request_data, + user_api_key_dict=user_api_key, + response=response, + ) + masked = await mask_guardrail.async_post_call_success_hook( + data=request_data, + user_api_key_dict=user_api_key, + response=response, + ) + assert masked.choices[0].message.content == expected, masked.choices[0].message.content + + async def _chunks(): + yield ModelResponseStream( + choices=[StreamingChoices(delta=Delta(content=echo, role="assistant"), index=0, finish_reason=None)] + ) + yield ModelResponseStream( + choices=[StreamingChoices(delta=Delta(content="", role="assistant"), index=0, finish_reason="stop")] + ) + + stream = restore_guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=user_api_key, + response=_chunks(), + request_data=request_data, + ) + stream = mask_guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=user_api_key, + response=stream, + request_data=request_data, + ) + collected = "" + async for chunk in stream: + if isinstance(chunk, ModelResponseStream): + for choice in chunk.choices: + if choice.delta and choice.delta.content: + collected += choice.delta.content + assert collected == expected, collected