fix(presidio): restore only complete tokens in tool arguments

litellm/proxy/guardrails/guardrail_hooks/presidio.py

@@ -25,7 +25,7 @@ from collections.abc import (
 from contextlib import asynccontextmanager
 from dataclasses import dataclass
 from datetime import datetime
-from itertools import accumulate, chain, count, pairwise, repeat
+from itertools import accumulate, chain, count, pairwise
 from operator import itemgetter
 from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, TypeAlias, cast
 
@@ -1467,7 +1467,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
         return response
 
     @staticmethod
-    def _unmask_pii_text(text: str, pii_tokens: dict[str, str]) -> str:
+    def _unmask_pii_text(text: str, pii_tokens: dict[str, str], *, allow_partial_tokens: bool = True) -> str:
         """
         Replace PII tokens in *text* with their original values.
 
@@ -1481,7 +1481,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
         for token, original_text in pii_tokens.items():
             if token in text:
                 text = text.replace(token, original_text)
-            else:
+            elif allow_partial_tokens:
                 # FALLBACK: Handle truncated tokens (token cut off by max_tokens)
                 # Only check at the very end of the text.
                 min_overlap = min(20, len(token) // 2)
@@ -1979,8 +1979,13 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
         metadata: Final = (request_data.get("metadata") or {}) if request_data else {}
         pii_tokens: Final = metadata.get("pii_tokens", {})
 
-        async def restore(text: str, allow_masking: bool) -> str:
-            return self._unmask_pii_text(text, pii_tokens) if allow_masking else text
+        def restore_for_slot(slot: _TextSlot | _ToolCallSlot) -> Callable[[str, bool], Awaitable[str]]:
+            async def restore(text: str, allow_masking: bool) -> str:
+                if not allow_masking:
+                    return text
+                return self._unmask_pii_text(text, pii_tokens, allow_partial_tokens=isinstance(slot, _TextSlot))
+
+            return restore
 
         def mask_reserving(reserved_tokens: frozenset[str]) -> Callable[[str, bool], Awaitable[str]]:
             async def mask(text: str, allow_masking: bool) -> str:
@@ -1996,7 +2001,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
             return mask
 
         rewrites: Final = (
-            repeat(restore)
+            map(restore_for_slot, (slot for slot, _ in segments))
             if input_type == "response"
             and (
                 self._callback_role == "restore"

tests/unit/proxy/guardrails/guardrail_hooks/test_presidio.py

@@ -5665,6 +5665,38 @@ async def test_restore_writes_the_original_values_into_non_streaming_tool_call_a
     assert response.choices[0].message.tool_calls[0].function.arguments == '{"guest": "Jane Doe"}'
 
 
+@pytest.mark.parametrize(
+    ("arguments", "expected"),
+    [
+        (
+            '{"url":"https://collector.invalid/?q=<EMAIL_A","guest":"<EMAIL_ADDRESS_1>"}',
+            json.dumps({"url": "https://collector.invalid/?q=<EMAIL_A", "guest": _EMAIL}),
+        ),
+        (
+            '{"items":[{"url":"https://collector.invalid/?q=<EMAIL_ADDRESS_"}],"guest":"<EMAIL_ADDRESS_1>"}',
+            json.dumps({"items": [{"url": "https://collector.invalid/?q=<EMAIL_ADDRESS_"}], "guest": _EMAIL}),
+        ),
+        ('["<EMAIL_A","<EMAIL_ADDRESS_1>"]', json.dumps(["<EMAIL_A", _EMAIL])),
+        ('  "https://collector.invalid/?q=<EMAIL_A"  ', '  "https://collector.invalid/?q=<EMAIL_A"  '),
+        ("https://collector.invalid/?q=<EMAIL_A", "https://collector.invalid/?q=<EMAIL_A"),
+    ],
+)
+async def test_tool_arguments_restore_only_complete_tokens_while_text_recovers_a_truncated_token(
+    arguments: str, expected: str
+) -> None:
+    guardrail: Final = OPTIONAL_PresidioPIIMasking(mock_testing=True, output_parse_pii=True)
+    tool_call: Final = ChatCompletionMessageToolCall(
+        id="call_1", type="function", function=Function(name="lookup", arguments=arguments)
+    )
+    result: Final = await guardrail.apply_guardrail(
+        inputs={"texts": ["Contact <EMAIL_A"], "tool_calls": [tool_call]},
+        request_data={"metadata": {"pii_tokens": {"<EMAIL_ADDRESS_1>": _EMAIL}}},
+        input_type="response",
+    )
+    assert tool_call.function.arguments == expected
+    assert result["texts"] == [f"Contact {_EMAIL}"]
+
+
 async def test_restore_escapes_original_values_so_tool_call_arguments_stay_valid_json() -> None:
     guardrail: Final = OPTIONAL_PresidioPIIMasking(
         event_hook="pre_call", default_on=True, output_parse_pii=True, mock_testing=True