From 212059cd118029f5f65df2883ad662f02730c6b9 Mon Sep 17 00:00:00 2001 From: Harshit28j Date: Tue, 10 Mar 2026 16:33:41 +0530 Subject: [PATCH] fix: presidio improvements --- .../guardrails/guardrail_hooks/presidio.py | 106 +++++- .../guardrail_hooks/test_presidio.py | 342 ++++++++++++++++++ 2 files changed, 435 insertions(+), 13 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index 8496a63ec31..507a0b89fc9 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -484,7 +484,11 @@ 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 @@ -515,9 +519,9 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): 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) @@ -902,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 ) @@ -917,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 @@ -945,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, @@ -1064,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. """ @@ -1081,6 +1144,10 @@ 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: return @@ -1134,6 +1201,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 @@ -1211,15 +1282,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 diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py index 6b3fdf2a81f..296faffa1a5 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py @@ -1843,3 +1843,345 @@ def test_unmask_graceful_degradation(): 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 at ", + "items": [ + { + "start": 5, + "end": 9, + "entity_type": "PERSON", + "text": "", + "operator": "replace", + }, + { + "start": 13, + "end": 25, + "entity_type": "PHONE_NUMBER", + "text": "", + "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": { + "": "John Smith", + "": "555-123-4567", + } + }, + } + + anthropic_response = { + "type": "message", + "id": "msg_123", + "model": "claude-3-haiku", + "role": "assistant", + "content": [ + { + "type": "text", + "text": "Hello , your number is .", + } + ], + "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": {"": "John"}}, + } + + anthropic_response = { + "type": "message", + "id": "msg_123", + "content": [ + {"type": "text", "text": "Hello "}, + { + "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": {"": "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": { + "": "John Smith", + "": "555-123-4567", + } + }, + } + + inputs = { + "texts": [ + "Hello , your number is .", + ] + } + + 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", "") + + 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 "" in result["texts"][0] + assert "John Smith" not in result["texts"][0]