Merge d54353db06 into ff2709f5c3

litellm/llms/anthropic/chat/guardrail_translation/handler.py

@@ -16,7 +16,7 @@ import json
 from collections.abc import Iterable, Mapping, MutableSequence, Sequence
 from copy import deepcopy
 from dataclasses import dataclass
-from itertools import chain, repeat
+from itertools import accumulate, chain, repeat
 from types import MappingProxyType
 from typing import TYPE_CHECKING, Any, Final, Protocol, TypeVar, cast, overload, runtime_checkable
 
@@ -519,6 +519,17 @@ def _rewritten_event(event: Mapping[str, object], rewrite_event: _SSEEventRewrit
     return {**event, rewrite.section: {**section, rewrite.field: rewrite.value}}
 
 
+def _tool_call_text_offsets(texts_before_messages: int, extracted: Sequence[ExtractedInput]) -> tuple[int, ...]:
+    texts_through_message: Final = tuple(
+        accumulate((len(one_message.scanned) for one_message in extracted), initial=texts_before_messages)
+    )[1:]
+    return tuple(
+        chain.from_iterable(
+            repeat(offset, len(one_message.tool_calls)) for one_message, offset in zip(extracted, texts_through_message)
+        )
+    )
+
+
 def _tool_call_shapes(tool_calls: Sequence[object]) -> tuple[_ToolCallShape, ...]:
     """The guardrail-visible shape of each tool call, whether the guardrail handed
     back the ``ChatCompletionMessageToolCall`` objects it was given or plain dicts."""
@@ -908,7 +919,14 @@ class AnthropicMessagesHandler(BaseTranslation):
 
         # Step 2: Apply guardrail to all texts and tool calls in batch
         if texts_to_check or tool_calls_to_check:
-            inputs: Final = GenericGuardrailAPIInputs(texts=texts_to_check)
+            inputs: Final = (
+                GenericGuardrailAPIInputs(
+                    texts=texts_to_check,
+                    tool_call_text_offsets=_tool_call_text_offsets(len(top_level_system_scanned), extracted),
+                )
+                if tool_calls_to_check
+                else GenericGuardrailAPIInputs(texts=texts_to_check)
+            )
             if images_to_check:
                 inputs["images"] = images_to_check
             if tool_calls_to_check:

litellm/llms/openai/chat/guardrail_translation/handler.py

@@ -17,6 +17,7 @@ This pattern can be replicated for other message formats (e.g., Anthropic).
 import json
 import time
 import uuid
+from bisect import bisect_right
 from collections.abc import Iterator, Mapping, Sequence
 from types import MappingProxyType
 from typing import TYPE_CHECKING, Any, Final, Union, cast
@@ -76,6 +77,13 @@ if TYPE_CHECKING:
     from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
 
 
+def _tool_call_text_offsets(
+    text_task_mappings: Sequence[tuple[int, int | None]], tool_call_task_mappings: Sequence[tuple[int, int]]
+) -> tuple[int, ...]:
+    text_message_indices: Final = tuple(msg_idx for msg_idx, _ in text_task_mappings)
+    return tuple(bisect_right(text_message_indices, msg_idx) for msg_idx, _ in tool_call_task_mappings)
+
+
 class OpenAIChatCompletionsHandler(BaseTranslation):
     """
     Handler for processing OpenAI chat completions messages with guardrails.
@@ -141,7 +149,14 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
 
         # Step 2: Apply guardrail to all texts and tool calls in batch
         if texts_to_check or tool_calls_to_check:
-            inputs: Final = GenericGuardrailAPIInputs(texts=texts_to_check)
+            inputs: Final = (
+                GenericGuardrailAPIInputs(
+                    texts=texts_to_check,
+                    tool_call_text_offsets=_tool_call_text_offsets(text_task_mappings, tool_call_task_mappings),
+                )
+                if tool_calls_to_check
+                else GenericGuardrailAPIInputs(texts=texts_to_check)
+            )
             if images_to_check:
                 inputs["images"] = images_to_check
             if tool_calls_to_check:

litellm/proxy/_lazy_openapi_snapshot.json

@@ -50261,6 +50261,13 @@
               "title": "Texts",
               "type": "array"
             },
+            "tool_call_text_offsets": {
+              "items": {
+                "type": "integer"
+              },
+              "title": "Tool Call Text Offsets",
+              "type": "array"
+            },
             "tool_calls": {
               "anyOf": [
                 {

litellm/proxy/guardrails/_content_utils.py

@@ -97,10 +97,14 @@ def _text_object(value: object) -> Mapping[str, object] | None:
         return None
 
 
+class DuplicateToolArgumentKeyError(ValueError):
+    pass
+
+
 def _tool_json_object(pairs: Sequence[tuple[str, JsonValue]]) -> JsonValue:
     value: Final = dict(pairs)
     if len(value) != len(pairs):
-        raise ValueError("Duplicate tool argument keys")
+        raise DuplicateToolArgumentKeyError("Duplicate tool argument keys")
     return value
 
 
@@ -111,7 +115,7 @@ def _finite_number(literal: str) -> float:
     return number
 
 
-def _tool_arguments(arguments: str) -> JsonValue:
+def parse_tool_arguments(arguments: str) -> JsonValue:
     parsed: Final = _tool_json(
         json.loads(
             arguments,
@@ -363,7 +367,7 @@ def walk_user_text(
         if not isinstance(arguments, str) or not arguments:
             return arguments
         try:
-            parsed: Final = _tool_arguments(arguments)
+            parsed: Final = parse_tool_arguments(arguments)
         except (ValueError, RecursionError):
             visited += 1  # rebind-ok: malformed arguments still undergo the text callback
             return visit(arguments)

litellm/proxy/guardrails/guardrail_hooks/deepkeep/deepkeep.py

@@ -253,8 +253,7 @@ class DeepKeepGuardrail(CustomGuardrail):
             getattr(logging_obj, "litellm_trace_id", None) if logging_obj else None,
             exc_info=error,
         )
-        return_inputs: Final[GenericGuardrailAPIInputs] = {}
-        return_inputs.update(inputs)
+        return_inputs: Final[GenericGuardrailAPIInputs] = {**inputs}
         return return_inputs
 
     def _handle_guardrail_request_error(

litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py

@@ -413,8 +413,7 @@ class GenericGuardrailAPI(CustomGuardrail):
             exc_info=error,
         )
         # Keep flow going - treat as action=NONE (no modifications)
-        return_inputs: Final[GenericGuardrailAPIInputs] = {}
-        return_inputs.update(inputs)
+        return_inputs: Final[GenericGuardrailAPIInputs] = {**inputs}
         return return_inputs
 
     def _build_request_headers(self) -> Mapping[str, str]:

litellm/proxy/guardrails/guardrail_hooks/presidio.py

@@ -12,12 +12,22 @@ import asyncio
 import json
 import re
 import threading
-from collections.abc import AsyncGenerator, AsyncIterable, AsyncIterator, Awaitable, Iterator, Sequence
+from collections.abc import (
+    AsyncGenerator,
+    AsyncIterable,
+    AsyncIterator,
+    Awaitable,
+    Callable,
+    Iterator,
+    Mapping,
+    Sequence,
+)
 from contextlib import asynccontextmanager
 from dataclasses import dataclass
 from datetime import datetime
-from itertools import accumulate, count, pairwise
-from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, cast
+from itertools import accumulate, chain, count, pairwise
+from operator import itemgetter
+from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, TypeAlias, cast
 
 import aiohttp
 from pydantic import ConfigDict, JsonValue, TypeAdapter, with_config
@@ -45,6 +55,7 @@ from litellm.integrations.custom_guardrail import (
 )
 from litellm.proxy._types import UserAPIKeyAuth
 from litellm.proxy.common_utils.sse_keepalive import split_complete_sse_frames
+from litellm.proxy.guardrails._content_utils import DuplicateToolArgumentKeyError, parse_tool_arguments
 from litellm.proxy.guardrails.anthropic_sse import (
     anthropic_sse_chunks_from_response,
     assemble_anthropic_sse_stream,
@@ -59,11 +70,12 @@ from litellm.types.guardrails import (
     PiiEntityType,
     PresidioPerRequestConfig,
 )
+from litellm.types.llms.openai import ChatCompletionToolCallChunk
 from litellm.types.proxy.guardrails.guardrail_hooks.presidio import (
     PresidioAnalyzeRequest,
     PresidioAnalyzeResponseItem,
 )
-from litellm.types.utils import GuardrailStatus, Message, StreamingChoices
+from litellm.types.utils import ChatCompletionMessageToolCall, Function, GuardrailStatus, Message, StreamingChoices
 from litellm.utils import (
     EmbeddingResponse,
     ImageResponse,
@@ -249,6 +261,138 @@ def _resolve_overlapping_spans(analyze_results: Sequence[PresidioAnalyzeResponse
     return tuple(_collapse_run(detections[first:last]) for first, last in pairwise(boundaries) if first < last)
 
 
+_LITERAL_TOKEN: Final = re.compile(r"<[^<>\s]+_\d+>")
+
+_ToolCall: TypeAlias = ChatCompletionToolCallChunk | ChatCompletionMessageToolCall
+
+
+@dataclass(frozen=True, slots=True)
+class _TextSlot:
+    index: int
+
+
+@dataclass(frozen=True, slots=True)
+class _ToolCallSlot:
+    index: int
+
+
+def _literal_tokens(text: str) -> frozenset[str]:
+    return frozenset(text[match.start() : match.end()] for match in _LITERAL_TOKEN.finditer(text))
+
+
+def _free_token(entity_type: str, numbers: Iterator[int], held: Mapping[str, str], reserved: frozenset[str]) -> str:
+    candidates: Final = (f"<{entity_type}_{number}>" for number in numbers)
+    return next(token for token in candidates if token not in held and token not in reserved)
+
+
+def _tool_call_arguments(tool_call: _ToolCall) -> str | None:
+    match tool_call:
+        case {"function": {"arguments": str() as arguments}}:
+            return arguments
+        case ChatCompletionMessageToolCall(function=Function(arguments=str() as arguments)):
+            return arguments
+        case _:
+            return None
+
+
+def _set_tool_call_arguments(tool_call: _ToolCall, arguments: str) -> None:
+    if isinstance(tool_call, ChatCompletionMessageToolCall):
+        function: Final = tool_call.function
+        function.arguments = arguments
+    else:
+        function_chunk: Final = tool_call["function"]
+        function_chunk["arguments"] = arguments
+
+
+def _parsed_arguments(arguments: str) -> JsonValue:
+    try:
+        return parse_tool_arguments(arguments)
+    except DuplicateToolArgumentKeyError as error:
+        raise GuardrailRaisedException(
+            message="Presidio cannot safely scan tool arguments with duplicate JSON object keys",
+            should_wrap_with_default_message=False,
+            blocked_content=True,
+        ) from error
+
+
+def _json_literal_tokens(value: JsonValue) -> frozenset[str]:
+    match value:
+        case str():
+            return _literal_tokens(value)
+        case list():
+            return frozenset[str]().union(*map(_json_literal_tokens, value))
+        case dict():
+            return frozenset[str]().union(*map(_literal_tokens, value), *map(_json_literal_tokens, value.values()))
+        case _:
+            return frozenset()
+
+
+def _segment_literal_tokens(segment: tuple[_TextSlot | _ToolCallSlot, str]) -> frozenset[str]:
+    slot, source = segment
+    if isinstance(slot, _TextSlot):
+        return _literal_tokens(source)
+    try:
+        parsed: Final = _parsed_arguments(source)
+    except ValueError:
+        return _literal_tokens(source)
+    return _json_literal_tokens(parsed)
+
+
+async def _rewritten_json_strings(value: JsonValue, rewrite: Callable[[str, bool], Awaitable[str]]) -> JsonValue:
+    match value:
+        case str():
+            return await rewrite(value, True)
+        case list():
+            return [await _rewritten_json_strings(item, rewrite) for item in value]
+        case dict():
+            return {
+                await rewrite(key, False): await _rewritten_json_strings(item, rewrite) for key, item in value.items()
+            }
+        case bool() | None:
+            return value
+        case int() | float():
+            await rewrite(str(value), False)
+            return value
+
+
+async def _rewritten_arguments(arguments: str, rewrite: Callable[[str, bool], Awaitable[str]]) -> str:
+    """
+    Rewrite each string of a tool call's JSON arguments on its own. An entity recognizer
+    finds a bare value that it misses inside JSON quoting, and a restored value with a
+    quote or a newline has to be escaped for the arguments to stay valid JSON.
+    """
+    try:
+        parsed: Final = _parsed_arguments(arguments)
+    except ValueError:
+        return await rewrite(arguments, True)
+    rewritten: Final = await _rewritten_json_strings(parsed, rewrite)
+    return arguments if rewritten == parsed else json.dumps(rewritten, ensure_ascii=False)
+
+
+def _conversation_segments(
+    texts: Sequence[str], arguments: Sequence[str | None], tool_call_text_offsets: Sequence[int]
+) -> Iterator[tuple[_TextSlot | _ToolCallSlot, str]]:
+    """
+    Every text and tool call argument string in the order the conversation holds them,
+    so masking an earlier message never depends on the turns appended after it.
+    """
+    offsets: Final = (
+        tool_call_text_offsets if len(tool_call_text_offsets) == len(arguments) else (len(texts),) * len(arguments)
+    )
+    keyed: Final = sorted(
+        chain(
+            (((index, 1), _TextSlot(index), text) for index, text in enumerate(texts)),
+            (
+                ((offset, 0), _ToolCallSlot(index), source)
+                for index, (offset, source) in enumerate(zip(offsets, arguments))
+                if source is not None
+            ),
+        ),
+        key=itemgetter(0),
+    )
+    return ((slot, source) for _, slot, source in keyed)
+
+
 class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
     user_api_key_cache = None
     ad_hoc_recognizers: list[str] | None = None
@@ -872,6 +1016,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
         analyze_results: Any,
         request_data: dict | None,
         masked_entity_count: dict[str, int],
+        reserved_tokens: frozenset[str] = frozenset(),
     ) -> str:
         # output_parse_pii is True — we need sequentially numbered
         # tokens and a pii_tokens mapping for later unmasking.
@@ -893,8 +1038,8 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
 
         runs: Final = _resolve_overlapping_spans(analyze_results)
         # Every message and content part numbers its runs from 1 into this one map, so a
-        # token another part already holds for a different value takes a spare number.
-        # No token in the map is numbered above its size, so spares start above both.
+        # token another part already holds for a different value, or one the request already
+        # carries as literal text, takes a spare number that is neither held nor reserved.
         spare_numbers: Final = count(max(len(pii_tokens), len(runs)) + 1)
 
         for seq, run in reversed(tuple(enumerate(runs, start=1))):
@@ -902,8 +1047,8 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
             numbered = f"<{run.entity_type}_{seq}>"
             replacement = (
                 numbered
-                if pii_tokens.get(numbered, original) == original
-                else f"<{run.entity_type}_{next(spare_numbers)}>"
+                if numbered not in reserved_tokens and pii_tokens.get(numbered, original) == original
+                else _free_token(run.entity_type, spare_numbers, pii_tokens, reserved_tokens)
             )
             pii_tokens[replacement] = original
             new_text = new_text[: run.start] + replacement + new_text[run.end :]
@@ -917,6 +1062,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
         output_parse_pii: bool,
         masked_entity_count: dict[str, int],
         request_data: dict | None = None,
+        reserved_tokens: frozenset[str] = frozenset(),
     ) -> str:
         """
         Send analysis results to the Presidio anonymizer endpoint to get redacted text
@@ -936,7 +1082,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
                 return self._finalize_presidio_anonymize_simple(redacted_text, masked_entity_count)
 
             return self._finalize_presidio_anonymize_numbered_tokens(
-                text, analyze_results, request_data, masked_entity_count
+                text, analyze_results, request_data, masked_entity_count, reserved_tokens
             )
         except Exception as e:
             # Sanitize exception to avoid leaking the original text (which may
@@ -987,7 +1133,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
         return filtered_results
 
     def raise_exception_if_blocked_entities_detected(
-        self, analyze_results: list[PresidioAnalyzeResponseItem] | _PresidioAnonymizeResponse
+        self, analyze_results: list[PresidioAnalyzeResponseItem] | _PresidioAnonymizeResponse, block_all: bool = False
     ) -> None:
         """
         Raise an exception if blocked entities are detected
@@ -1005,7 +1151,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
 
             if entity_type:
                 # Check if entity_type is in config (supports both enum and string)
-                if entity_type in self.pii_entities_config and self.pii_entities_config[entity_type] == PiiAction.BLOCK:
+                if block_all or self.pii_entities_config.get(entity_type) == PiiAction.BLOCK:
                     raise BlockedPiiEntityError(
                         entity_type=entity_type,
                         guardrail_name=self.guardrail_name,
@@ -1017,6 +1163,8 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
         output_parse_pii: bool,
         presidio_config: PresidioPerRequestConfig | None,
         request_data: dict,
+        reserved_tokens: frozenset[str] = frozenset(),
+        allow_masking: bool = True,
     ) -> str:
         """
         Calls Presidio Analyze + Anonymize endpoints for PII Analysis + Masking
@@ -1045,7 +1193,11 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
                 ####################################################
                 # Blocked Entities check
                 ####################################################
-                self.raise_exception_if_blocked_entities_detected(analyze_results=analyze_results)
+                self.raise_exception_if_blocked_entities_detected(
+                    analyze_results=analyze_results, block_all=not allow_masking
+                )
+                if not allow_masking:
+                    return text
 
                 # Then anonymize the text using the analysis results
                 anonymized_text: Final = await self.anonymize_text(
@@ -1054,6 +1206,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
                     output_parse_pii=output_parse_pii,
                     masked_entity_count=masked_entity_count,
                     request_data=request_data,
+                    reserved_tokens=reserved_tokens,
                 )
                 return anonymized_text
             return redacted_text["text"]
@@ -1314,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.
 
@@ -1328,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)
@@ -1816,27 +1969,58 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
             2. When Testing the guardrail with some text, this function will be called with the input text and returns a text after applying the guardrail
         """
         texts: Final = inputs.get("texts", [])
+        tool_calls: Final = inputs.get("tool_calls") or ()
+        segments: Final = tuple(
+            _conversation_segments(
+                texts, tuple(map(_tool_call_arguments, tool_calls)), inputs.get("tool_call_text_offsets", ())
+            )
+        )
 
         metadata: Final = (request_data.get("metadata") or {}) if request_data else {}
         pii_tokens: Final = metadata.get("pii_tokens", {})
 
-        new_texts: Final = []
-        if input_type == "response" and (
-            self._callback_role == "restore"
-            or (self._callback_role is None and not self.apply_to_output 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(
+        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:
+                return await self.check_pii(
                     text=text,
                     output_parse_pii=self.output_parse_pii,
                     presidio_config=None,
                     request_data=request_data or {},
+                    reserved_tokens=reserved_tokens,
+                    allow_masking=allow_masking,
                 )
-                new_texts.append(modified_text)
-        inputs["texts"] = new_texts
+
+            return mask
+
+        rewrites: Final = (
+            map(restore_for_slot, (slot for slot, _ in segments))
+            if input_type == "response"
+            and (
+                self._callback_role == "restore"
+                or (self._callback_role is None and not self.apply_to_output and pii_tokens)
+            )
+            else map(
+                mask_reserving,
+                accumulate(map(_segment_literal_tokens, segments), frozenset[str].union),
+            )
+        )
+        rewritten: Final = [
+            await (_rewritten_arguments(source, rewrite) if isinstance(slot, _ToolCallSlot) else rewrite(source, True))
+            for (slot, source), rewrite in zip(segments, rewrites)
+        ]
+        by_slot: Final = dict(zip((slot for slot, _ in segments), rewritten))
+        inputs["texts"] = [by_slot[_TextSlot(index)] for index in range(len(texts))]
+        for index, tool_call in enumerate(tool_calls):
+            if _ToolCallSlot(index) in by_slot:
+                _set_tool_call_arguments(tool_call, by_slot[_ToolCallSlot(index)])
         return inputs
 
     def update_in_memory_litellm_params(self, litellm_params: LitellmParams) -> None:

litellm/proxy/guardrails/guardrail_hooks/straiker/straiker.py

@@ -1213,8 +1213,7 @@ class StraikerGuardrail(CustomGuardrail):
         inputs: GenericGuardrailAPIInputs,
         parsed: StraikerWebhookResponse,
     ) -> GenericGuardrailAPIInputs:
-        return_inputs: Final[GenericGuardrailAPIInputs] = {}
-        return_inputs.update(inputs)
+        return_inputs: Final[GenericGuardrailAPIInputs] = {**inputs}
         if parsed.texts is not None:
             return_inputs["texts"] = parsed.texts
         return return_inputs

litellm/types/utils.py

@@ -5361,6 +5361,7 @@ class GenericGuardrailAPIInputs(TypedDict, total=False):
     images: list[str]  # extracted images from the LLM response - for image guardrails
     tools: list[ChatCompletionToolParam]  # tools sent to the LLM
     tool_calls: list[ChatCompletionToolCallChunk] | list[ChatCompletionMessageToolCall]  # tool calls sent from the LLM
+    tool_call_text_offsets: ReadOnly[tuple[int, ...]]  # per tool call, how many texts precede it in the conversation
     structured_messages: list[
         AllMessageValues
     ]  # structured messages sent to the LLM - indicates if text is from system or user

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

@@ -26,14 +26,25 @@ import litellm
 from litellm import mock_completion
 from litellm.caching.caching import DualCache
 from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException
+from litellm.llms.anthropic.chat.guardrail_translation.handler import AnthropicMessagesHandler
+from litellm.llms.openai.chat.guardrail_translation.handler import OpenAIChatCompletionsHandler
 from litellm.proxy._types import UserAPIKeyAuth
 from litellm.proxy.guardrails.guardrail_hooks.presidio import (
     OPTIONAL_PresidioPIIMasking,
     PresidioPerRequestConfig,
 )
+from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import UnifiedLLMGuardrails
 from litellm.types.guardrails import LitellmParams, PiiAction, PiiEntityType
 from litellm.types.proxy.guardrails.guardrail_hooks.presidio import PresidioAnalyzeRequest, PresidioAnalyzeResponseItem
-from litellm.types.utils import Choices, Delta, Message, ModelResponse, StreamingChoices
+from litellm.types.utils import (
+    ChatCompletionMessageToolCall,
+    Choices,
+    Delta,
+    Function,
+    Message,
+    ModelResponse,
+    StreamingChoices,
+)
 
 
 def _make_mock_session_iterator(json_response, status=200, content_type="application/json", text_response=""):
@@ -696,7 +707,9 @@ async def test_presidio_sets_guardrail_information_in_request_data():
         "metadata": {},
     }
 
-    async def mock_check_pii(text, output_parse_pii, presidio_config, request_data):
+    async def mock_check_pii(
+        text, output_parse_pii, presidio_config, request_data, reserved_tokens=frozenset(), *, allow_masking: bool = True
+    ):
         assert request_data is not None
 
         presidio.add_standard_logging_guardrail_information_to_request_data(
@@ -754,7 +767,9 @@ async def test_request_data_flows_to_apply_guardrail():
         "metadata": {},
     }
 
-    async def mock_check_pii(text, output_parse_pii, presidio_config, request_data):
+    async def mock_check_pii(
+        text, output_parse_pii, presidio_config, request_data, reserved_tokens=frozenset(), *, allow_masking: bool = True
+    ):
         assert request_data is not None, "request_data should be passed to check_pii"
         assert "metadata" in request_data, "request_data should have metadata"
 
@@ -2336,7 +2351,9 @@ async def test_standalone_scans_without_restoration_tokens(input_type: Literal["
         mock_testing=True,
     )
 
-    async def mock_check_pii(text, output_parse_pii, presidio_config, request_data):
+    async def mock_check_pii(
+        text, output_parse_pii, presidio_config, request_data, reserved_tokens=frozenset(), *, allow_masking: bool = True
+    ):
         return text.replace("John Smith", "<PERSON>")
 
     guardrail.check_pii = mock_check_pii
@@ -5222,7 +5239,11 @@ def _card_and_email_spans(text: str) -> tuple[PresidioAnalyzeResponseItem, ...]:
             score=1.0,
             analysis_explanation=None,
         )
-        for entity_type, value in (("CREDIT_CARD", _CARD_NUMBER), ("EMAIL_ADDRESS", _EMAIL))
+        for entity_type, value in (
+            ("CREDIT_CARD", _CARD_NUMBER),
+            ("CREDIT_CARD", _CARD_NUMBER.replace("-", "")),
+            ("EMAIL_ADDRESS", _EMAIL),
+        )
         if value in text
     )
 
@@ -5310,9 +5331,15 @@ def _drain(analyzed: asyncio.Queue[PresidioAnalyzeRequest]) -> tuple[PresidioAna
 
 def _guardrail_with_in_memory_presidio(
     analyzed: asyncio.Queue[PresidioAnalyzeRequest],
+    *,
+    output_parse_pii: bool = False,
+    action: PiiAction | None = None,
 ) -> OPTIONAL_PresidioPIIMasking:
     guardrail: Final = OPTIONAL_PresidioPIIMasking(
-        pii_entities_config=_BLOCK_CARD_MASK_EMAIL,
+        output_parse_pii=output_parse_pii,
+        pii_entities_config=(
+            _BLOCK_CARD_MASK_EMAIL if action is None else {entity: action for entity in _BLOCK_CARD_MASK_EMAIL}
+        ),
         presidio_analyzer_api_base="http://presidio-analyzer.test/",
         presidio_anonymizer_api_base="http://presidio-anonymizer.test/",
     )
@@ -5438,3 +5465,442 @@ 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 _function_tool_call(call_id: str, arguments: str) -> dict:
+    return {"id": call_id, "type": "function", "function": {"name": "lookup", "arguments": arguments}}
+
+
+def _tool_call_response(arguments: str) -> ModelResponse:
+    return ModelResponse(
+        choices=[
+            Choices(
+                message=Message(
+                    role="assistant",
+                    content=None,
+                    tool_calls=[
+                        ChatCompletionMessageToolCall(
+                            id="call_1", type="function", function=Function(name="lookup", arguments=arguments)
+                        )
+                    ],
+                )
+            )
+        ]
+    )
+
+
+_UNQUOTED_PERSON_NAME: Final = re.compile(r'(?<!")\b[A-Z][a-z]+ [A-Z][a-z]+\b(?!")')
+
+
+async def _analyze_names_outside_json_quotes(request: web.Request) -> web.Response:
+    payload = await request.json()
+    return web.json_response(
+        [
+            {"entity_type": "PERSON", "start": match.start(), "end": match.end(), "score": 0.85}
+            for match in _UNQUOTED_PERSON_NAME.finditer(payload["text"])
+        ]
+    )
+
+
+def _json_blind_presidio_app() -> web.Application:
+    app = web.Application()
+    app.router.add_post("/analyze", _analyze_names_outside_json_quotes)
+    app.router.add_post("/anonymize", _fake_anonymize)
+    return app
+
+
+@asynccontextmanager
+async def _reversible_presidio_guardrail(
+    presidio_app: web.Application | None = None,
+) -> AsyncIterator[OPTIONAL_PresidioPIIMasking]:
+    async with TestServer(presidio_app or _fake_presidio_app()) as server:
+        guardrail: Final = OPTIONAL_PresidioPIIMasking(
+            presidio_analyzer_api_base=str(server.make_url("/")),
+            presidio_anonymizer_api_base=str(server.make_url("/")),
+            output_parse_pii=True,
+            pii_entities_config={PiiEntityType.PERSON: PiiAction.MASK},
+        )
+        yield guardrail
+        await guardrail._close_http_session()
+
+
+async def _masked_chat_request(guardrail: OPTIONAL_PresidioPIIMasking, messages: Sequence[dict]) -> dict:
+    data: Final = {"model": "claude-fable-5-1", "messages": copy.deepcopy(list(messages))}
+    await OpenAIChatCompletionsHandler().process_input_messages(data=data, guardrail_to_apply=guardrail)
+    return data
+
+
+async def _masked_anthropic_request(guardrail: OPTIONAL_PresidioPIIMasking, messages: Sequence[dict]) -> dict:
+    data: Final = {"model": "claude-fable-5-1", "max_tokens": 64, "messages": copy.deepcopy(list(messages))}
+    await AnthropicMessagesHandler().process_input_messages(data=data, guardrail_to_apply=guardrail)
+    return data
+
+
+@pytest.mark.parametrize("action", [PiiAction.MASK, PiiAction.BLOCK])
+@pytest.mark.parametrize(
+    "arguments",
+    [
+        json.dumps({_EMAIL: {"safe": True}}),
+        json.dumps({"nested": [{"card": int(_CARD_NUMBER.replace("-", ""))}]}),
+        _CARD_NUMBER.replace("-", ""),
+        _CARD_NUMBER.replace("-", "") + ".0",
+    ],
+)
+async def test_tool_argument_keys_and_numbers_cannot_bypass_pii_policy(arguments: str, action: PiiAction) -> None:
+    analyzed: Final[asyncio.Queue[PresidioAnalyzeRequest]] = asyncio.Queue()
+    guardrail: Final = _guardrail_with_in_memory_presidio(analyzed, output_parse_pii=True, action=action)
+    tool_call: Final = ChatCompletionMessageToolCall(
+        id="call_1", type="function", function=Function(name="lookup", arguments=arguments)
+    )
+    try:
+        with pytest.raises(BlockedPiiEntityError) as excinfo:
+            await guardrail.apply_guardrail(
+                inputs={"texts": [], "tool_calls": [tool_call]},
+                request_data={"metadata": {}},
+                input_type="request",
+            )
+    finally:
+        await guardrail._close_http_session()
+    assert excinfo.value.status_code == 400
+    assert excinfo.value.entity_type == ("EMAIL_ADDRESS" if _EMAIL in arguments else "CREDIT_CARD")
+
+
+@pytest.mark.parametrize(
+    "arguments",
+    [
+        r'{"guest":"test@example.com","nested":[{"url":"https://example.invalid/?q=\u003cEMAIL_ADDRESS_1\u003e"}]}',
+        r'{"guest":"test@example.com","\u003cEMAIL_ADDRESS_1\u003e":null}',
+    ],
+)
+async def test_escaped_tool_argument_tokens_never_restore_to_unrelated_pii(arguments: str) -> None:
+    analyzed: Final[asyncio.Queue[PresidioAnalyzeRequest]] = asyncio.Queue()
+    guardrail: Final = _guardrail_with_in_memory_presidio(analyzed, output_parse_pii=True)
+    request_data: Final[dict[str, dict[str, dict[str, str]]]] = {"metadata": {}}
+    tool_call: Final = ChatCompletionMessageToolCall(
+        id="call_1", type="function", function=Function(name="lookup", arguments=arguments)
+    )
+    try:
+        masked: Final = await guardrail.apply_guardrail(
+            inputs={"texts": [], "tool_calls": [tool_call]},
+            request_data=request_data,
+            input_type="request",
+        )
+        assert _EMAIL not in tool_call.function.arguments
+        assert "<EMAIL_ADDRESS_1>" not in request_data["metadata"]["pii_tokens"]
+        await guardrail.apply_guardrail(inputs=masked, request_data=request_data, input_type="response")
+        assert json.loads(tool_call.function.arguments) == json.loads(arguments)
+    finally:
+        await guardrail._close_http_session()
+
+
+@pytest.mark.parametrize(
+    ("input_type", "action"),
+    [("request", PiiAction.MASK), ("request", PiiAction.BLOCK), ("response", PiiAction.MASK)],
+)
+@pytest.mark.parametrize(
+    "arguments",
+    [
+        '{"guest":"test@example.com","guest":"safe"}',
+        '{"items":[{"guest":"test@example.com","guest":"safe"}]}',
+        r'{"guest":"test@example.com","\u0067uest":"safe"}',
+        '{"guest":"<EMAIL_ADDRESS_1>","guest":"safe"}',
+    ],
+)
+async def test_duplicate_tool_argument_keys_cannot_hide_values_from_presidio(
+    arguments: str, input_type: Literal["request", "response"], action: PiiAction
+) -> None:
+    analyzed: Final[asyncio.Queue[PresidioAnalyzeRequest]] = asyncio.Queue()
+    guardrail: Final = _guardrail_with_in_memory_presidio(analyzed, output_parse_pii=True, action=action)
+    tool_call: Final = ChatCompletionMessageToolCall(
+        id="call_1", type="function", function=Function(name="lookup", arguments=arguments)
+    )
+    try:
+        with pytest.raises(GuardrailRaisedException, match="duplicate JSON object keys") as excinfo:
+            await guardrail.apply_guardrail(
+                inputs={"texts": [], "tool_calls": [tool_call]},
+                request_data={"metadata": {"pii_tokens": {"<EMAIL_ADDRESS_1>": _EMAIL}}},
+                input_type=input_type,
+            )
+        assert excinfo.value.status_code == 400
+        assert excinfo.value.blocked_content is True
+        assert _EMAIL not in str(excinfo.value)
+    finally:
+        await guardrail._close_http_session()
+
+
+@pytest.mark.parametrize(
+    "arguments",
+    [
+        '{ "count": 42, "fraction": 1.25, "enabled": true, "absent": null, "items": [false, 0] }',
+        '{ "left": {"guest": "safe"}, "right": {"guest": "safe"} }',
+    ],
+)
+async def test_safe_tool_argument_keys_and_scalar_types_keep_original_bytes(arguments: str) -> None:
+    analyzed: Final[asyncio.Queue[PresidioAnalyzeRequest]] = asyncio.Queue()
+    guardrail: Final = _guardrail_with_in_memory_presidio(analyzed, output_parse_pii=True)
+    tool_call: Final = ChatCompletionMessageToolCall(
+        id="call_1", type="function", function=Function(name="lookup", arguments=arguments)
+    )
+    try:
+        await guardrail.apply_guardrail(
+            inputs={"texts": [], "tool_calls": [tool_call]},
+            request_data={"metadata": {}},
+            input_type="request",
+        )
+        assert tool_call.function.arguments == arguments
+    finally:
+        await guardrail._close_http_session()
+
+
+async def test_restore_writes_the_original_values_into_non_streaming_tool_call_arguments() -> None:
+    guardrail: Final = OPTIONAL_PresidioPIIMasking(
+        event_hook="pre_call", default_on=True, output_parse_pii=True, mock_testing=True
+    )
+    response: Final = _tool_call_response('{"guest": "<PERSON_1>"}')
+    data: Final = {"metadata": {"pii_tokens": {"<PERSON_1>": "Jane Doe"}}, "guardrail_to_apply": guardrail}
+
+    await UnifiedLLMGuardrails().async_post_call_success_hook(
+        data, UserAPIKeyAuth(request_route="/v1/chat/completions"), response
+    )
+
+    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
+    )
+    original: Final = 'Apt "B"\nFloor 2'
+    response: Final = _tool_call_response('{"address": "<LOCATION_1>"}')
+    data: Final = {"metadata": {"pii_tokens": {"<LOCATION_1>": original}}, "guardrail_to_apply": guardrail}
+
+    await UnifiedLLMGuardrails().async_post_call_success_hook(
+        data, UserAPIKeyAuth(request_route="/v1/chat/completions"), response
+    )
+
+    assert json.loads(response.choices[0].message.tool_calls[0].function.arguments) == {"address": original}
+
+
+async def test_names_inside_json_tool_call_arguments_are_masked_even_when_the_analyzer_needs_bare_text() -> None:
+    async with _reversible_presidio_guardrail(_json_blind_presidio_app()) as guardrail:
+        data: Final = await _masked_chat_request(
+            guardrail,
+            [{"role": "assistant", "content": None, "tool_calls": [_function_tool_call("call_1", '{"guest": "Jane Doe"}')]}],
+        )
+
+    masked: Final = json.loads(data["messages"][0]["tool_calls"][0]["function"]["arguments"])
+    pii_tokens: Final = data.get("metadata", {}).get("pii_tokens", {})
+    assert pii_tokens.get(masked["guest"]) == "Jane Doe", f"masked arguments {masked} against token map {pii_tokens}"
+
+
+async def test_tool_call_arguments_without_pii_reach_the_model_byte_for_byte() -> None:
+    arguments: Final = '{"city":"Boston","days":[1,2]}'
+    async with _reversible_presidio_guardrail() as guardrail:
+        data: Final = await _masked_chat_request(
+            guardrail,
+            [
+                {"role": "user", "content": "Plan a trip for Jane Doe."},
+                {"role": "assistant", "content": None, "tool_calls": [_function_tool_call("call_1", arguments)]},
+            ],
+        )
+
+    assert data["messages"][0]["content"] != "Plan a trip for Jane Doe.", data["messages"][0]
+    assert data["messages"][1]["tool_calls"][0]["function"]["arguments"] == arguments
+
+
+@pytest.mark.parametrize(
+    "arguments",
+    [
+        pytest.param("call Jane Doe back", id="not-json"),
+        pytest.param('"Jane Doe"', id="json-string"),
+    ],
+)
+async def test_tool_call_arguments_that_are_not_a_json_object_are_still_masked_and_restored(arguments: str) -> None:
+    async with _reversible_presidio_guardrail() as guardrail:
+        data: Final = await _masked_chat_request(
+            guardrail,
+            [{"role": "assistant", "content": None, "tool_calls": [_function_tool_call("call_1", arguments)]}],
+        )
+
+    masked: Final = data["messages"][0]["tool_calls"][0]["function"]["arguments"]
+    assert "Jane Doe" not in masked, masked
+    assert OPTIONAL_PresidioPIIMasking._unmask_pii_text(masked, data["metadata"]["pii_tokens"]) == arguments, masked
+
+
+async def test_a_tool_call_without_string_arguments_passes_through_beside_masked_text() -> None:
+    tool_call: Final = {"id": "call_1", "type": "function", "function": {"name": "lookup"}}
+    async with _reversible_presidio_guardrail() as guardrail:
+        data: Final = await _masked_chat_request(
+            guardrail,
+            [
+                {"role": "user", "content": "Book for Jane Doe."},
+                {"role": "assistant", "content": None, "tool_calls": [tool_call]},
+            ],
+        )
+
+    assert "Jane Doe" not in data["messages"][0]["content"], data["messages"][0]
+    assert data["messages"][1]["tool_calls"] == [tool_call]
+
+
+async def test_apply_guardrail_restores_the_arguments_of_tool_call_objects() -> None:
+    guardrail: Final = OPTIONAL_PresidioPIIMasking(
+        event_hook="pre_call", default_on=True, output_parse_pii=True, mock_testing=True
+    )
+    tool_call: Final = ChatCompletionMessageToolCall(
+        id="call_1", type="function", function=Function(name="lookup", arguments='{"guest": "<PERSON_1>"}')
+    )
+
+    await guardrail.apply_guardrail(
+        inputs={"texts": [], "tool_calls": [tool_call]},
+        request_data={"metadata": {"pii_tokens": {"<PERSON_1>": "Jane Doe"}}},
+        input_type="response",
+    )
+
+    assert tool_call.function.arguments == '{"guest": "Jane Doe"}'
+
+
+async def test_output_masking_masks_non_streaming_tool_call_arguments() -> None:
+    async with TestServer(_fake_presidio_app()) as server:
+        guardrail: Final = OPTIONAL_PresidioPIIMasking(
+            presidio_analyzer_api_base=str(server.make_url("/")),
+            presidio_anonymizer_api_base=str(server.make_url("/")),
+            apply_to_output=True,
+            event_hook="post_call",
+            default_on=True,
+            pii_entities_config={PiiEntityType.PERSON: PiiAction.MASK},
+        )
+        response: Final = _tool_call_response('{"guest": "Jane Doe"}')
+        await UnifiedLLMGuardrails().async_post_call_success_hook(
+            {"metadata": {}, "guardrail_to_apply": guardrail},
+            UserAPIKeyAuth(request_route="/v1/chat/completions"),
+            response,
+        )
+        await guardrail._close_http_session()
+
+    assert response.choices[0].message.tool_calls[0].function.arguments == '{"guest": "<PERSON>"}'
+
+
+async def test_history_tool_call_arguments_reach_the_model_masked_and_restore_to_the_original() -> None:
+    messages: Final = [
+        {"role": "user", "content": "Book a table for Jane Doe."},
+        {
+            "role": "assistant",
+            "content": None,
+            "tool_calls": [_function_tool_call("call_1", '{"guest": "Jane Doe", "host": "John Smith"}')],
+        },
+        {"role": "tool", "tool_call_id": "call_1", "content": "Booked."},
+    ]
+    async with _reversible_presidio_guardrail() as guardrail:
+        data: Final = await _masked_chat_request(guardrail, messages)
+
+    masked: Final = json.loads(data["messages"][1]["tool_calls"][0]["function"]["arguments"])
+    pii_tokens: Final = data["metadata"]["pii_tokens"]
+    assert {key: pii_tokens.get(token) for key, token in masked.items()} == {
+        "guest": "Jane Doe",
+        "host": "John Smith",
+    }, f"masked arguments {masked} against token map {pii_tokens}"
+
+
+async def test_masking_history_tool_call_arguments_replays_a_byte_identical_prefix_across_turns() -> None:
+    earlier: Final = [
+        {"role": "user", "content": "My name is John Smith."},
+        {
+            "role": "assistant",
+            "content": None,
+            "tool_calls": [_function_tool_call("call_1", '{"name": "Alice Brown"}')],
+        },
+        {"role": "tool", "tool_call_id": "call_1", "content": "Found one match."},
+    ]
+    later: Final = [
+        *earlier,
+        {"role": "assistant", "content": "Noted."},
+        {"role": "user", "content": "Now add Bob Jones."},
+    ]
+    async with _reversible_presidio_guardrail() as guardrail:
+        first: Final = await _masked_chat_request(guardrail, earlier)
+        second: Final = await _masked_chat_request(guardrail, later)
+
+    assert "Alice Brown" not in json.dumps(first["messages"]), first["messages"]
+    assert second["messages"][: len(earlier)] == first["messages"]
+
+
+async def test_masking_anthropic_tool_use_input_replays_a_byte_identical_prefix_across_turns() -> None:
+    earlier: Final = [
+        {"role": "user", "content": "My name is John Smith."},
+        {
+            "role": "assistant",
+            "content": [{"type": "tool_use", "id": "toolu_1", "name": "lookup", "input": {"name": "Alice Brown"}}],
+        },
+        {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "toolu_1", "content": "Found one match."}]},
+    ]
+    later: Final = [
+        *earlier,
+        {"role": "assistant", "content": "Noted."},
+        {"role": "user", "content": "Now add Bob Jones."},
+    ]
+    async with _reversible_presidio_guardrail() as guardrail:
+        first: Final = await _masked_anthropic_request(guardrail, earlier)
+        second: Final = await _masked_anthropic_request(guardrail, later)
+
+    assert "Alice Brown" not in json.dumps(first["messages"]), first["messages"]
+    assert second["messages"][: len(earlier)] == first["messages"]
+
+
+@pytest.mark.parametrize(
+    ("contents", "literals"),
+    [
+        pytest.param(
+            ("The guest was <PERSON_1>.", "Book for Jane Doe."), ("<PERSON_1>",), id="literal-in-an-earlier-message"
+        ),
+        pytest.param(("Ask <PERSON_1> about Jane Doe.",), ("<PERSON_1>",), id="literal-beside-the-value"),
+        pytest.param(
+            ("Seen <PERSON_1> and <PERSON_2>.", "Book for Jane Doe.", "Add John Smith."),
+            ("<PERSON_1>", "<PERSON_2>"),
+            id="spare-skips-reserved-and-held-tokens",
+        ),
+    ],
+)
+async def test_a_token_the_request_already_carries_is_never_assigned_to_another_value(
+    contents: tuple[str, ...], literals: tuple[str, ...]
+) -> None:
+    async with _reversible_presidio_guardrail() as guardrail:
+        data: Final = await _masked_chat_request(
+            guardrail, [{"role": "user", "content": content} for content in contents]
+        )
+
+    masked: Final = tuple(message["content"] for message in data["messages"])
+    pii_tokens: Final = data["metadata"]["pii_tokens"]
+    assert "Jane Doe" not in " ".join(masked), masked
+    assert set(literals).isdisjoint(pii_tokens), pii_tokens
+    assert tuple(OPTIONAL_PresidioPIIMasking._unmask_pii_text(text, pii_tokens) for text in masked) == contents, masked

ui/litellm-dashboard/src/lib/http/schema.d.ts

@@ -34185,6 +34185,8 @@ export interface components {
             structured_messages?: (components["schemas"]["ChatCompletionUserMessage"] | components["schemas"]["ChatCompletionAssistantMessage"] | components["schemas"]["ChatCompletionToolMessage"] | components["schemas"]["ChatCompletionSystemMessage"] | components["schemas"]["ChatCompletionFunctionMessage"] | components["schemas"]["ChatCompletionDeveloperMessage"])[];
             /** Texts */
             texts?: string[];
+            /** Tool Call Text Offsets */
+            tool_call_text_offsets?: number[];
             /** Tool Calls */
             tool_calls?: components["schemas"]["ChatCompletionToolCallChunk"][] | components["schemas"]["ChatCompletionMessageToolCall"][];
             /** Tools */