From c4410080fc5312d2a8978080feda2de74ef9559a Mon Sep 17 00:00:00 2001 From: lawrence Date: Fri, 23 Oct 2026 22:41:32 +0800 Subject: [PATCH 1/6] fix(presidio): mask and restore tool call arguments in conversation order apply_guardrail only rewrote inputs["texts"], so on the unified guardrail path the assistant tool call arguments in a request reached the model unmasked, and placeholders in a non-streaming tool call came back to the client unrestored. Tool call arguments now go through the same masking and restoring as text, which the native hooks already did Masking walks texts and tool calls in the order the conversation holds them, using a tool_call_text_offsets input that the chat completions and Anthropic messages handlers fill in, so an earlier message masks the same way on every turn and the replayed prefix stays byte-stable A placeholder-shaped string the request already carries is never handed to a new value, so restoring the response cannot turn it into someone else's value. Only the same and earlier texts are reserved, which keeps earlier messages stable when later turns add placeholders Fixes #31950 --- .../chat/guardrail_translation/handler.py | 22 +- .../chat/guardrail_translation/handler.py | 17 +- litellm/proxy/_lazy_openapi_snapshot.json | 7 + .../guardrail_hooks/deepkeep/deepkeep.py | 3 +- .../generic_guardrail_api.py | 3 +- .../guardrails/guardrail_hooks/presidio.py | 131 +++++++++-- .../guardrail_hooks/straiker/straiker.py | 3 +- litellm/types/utils.py | 1 + .../guardrail_hooks/test_presidio.py | 214 +++++++++++++++++- ui/litellm-dashboard/src/lib/http/schema.d.ts | 2 + 10 files changed, 367 insertions(+), 36 deletions(-) diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index 4c6bded68ea..d849fbcb0d9 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/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, cast, overload, runtime_checkable @@ -336,6 +336,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.""" @@ -671,7 +682,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: diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index 3070bedd30d..0e9b1ca9bb0 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/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: diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 4572e59944b..9e10ecaa94b 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -48466,6 +48466,13 @@ "title": "Texts", "type": "array" }, + "tool_call_text_offsets": { + "items": { + "type": "integer" + }, + "title": "Tool Call Text Offsets", + "type": "array" + }, "tool_calls": { "anyOf": [ { diff --git a/litellm/proxy/guardrails/guardrail_hooks/deepkeep/deepkeep.py b/litellm/proxy/guardrails/guardrail_hooks/deepkeep/deepkeep.py index 23803b636f2..37b61cb8695 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/deepkeep/deepkeep.py +++ b/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( diff --git a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py index 199da7df30d..82e9c49d348 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py +++ b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py @@ -405,8 +405,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]: diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index cadc9ca689b..cf5a9913deb 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -12,12 +12,13 @@ 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, 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 @@ -58,11 +59,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, @@ -236,6 +238,74 @@ 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: + match tool_call: + case {"function": function}: + function["arguments"] = arguments + case ChatCompletionMessageToolCall(function=Function() as function): + function.arguments = arguments + case _: + pass + + +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 @@ -859,6 +929,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. @@ -880,8 +951,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))): @@ -889,8 +960,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 :] @@ -904,6 +975,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 @@ -923,7 +995,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 @@ -1004,6 +1076,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): output_parse_pii: bool, presidio_config: PresidioPerRequestConfig | None, request_data: dict, + reserved_tokens: frozenset[str] = frozenset(), ) -> str: """ Calls Presidio Analyze + Anonymize endpoints for PII Analysis + Masking @@ -1041,6 +1114,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"] @@ -1786,27 +1860,40 @@ 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( - text=text, + reserved: Final = accumulate(map(_literal_tokens, (source for _, source in segments)), frozenset[str].union) + rewritten: Final = ( + [self._unmask_pii_text(source, pii_tokens) for _, source 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 [ + await self.check_pii( + text=source, output_parse_pii=self.output_parse_pii, presidio_config=None, request_data=request_data or {}, + reserved_tokens=reserved_tokens, ) - new_texts.append(modified_text) - inputs["texts"] = new_texts + for (_, source), reserved_tokens in zip(segments, reserved) + ] + ) + 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: diff --git a/litellm/proxy/guardrails/guardrail_hooks/straiker/straiker.py b/litellm/proxy/guardrails/guardrail_hooks/straiker/straiker.py index f2f3ec201f5..9c45873421e 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/straiker/straiker.py +++ b/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 diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 5939a350970..0ba90ebad3e 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -4959,6 +4959,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 diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_presidio.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_presidio.py index bd1e9ae4068..95f1eaea0e8 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_presidio.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_presidio.py @@ -8,7 +8,7 @@ import copy import json import os import re -from collections.abc import Iterable, Sequence +from collections.abc import AsyncIterator, Iterable, Sequence from concurrent.futures import ThreadPoolExecutor from contextlib import asynccontextmanager from typing import Final, Literal @@ -26,15 +26,26 @@ from typing_extensions import ReadOnly, TypedDict import litellm from litellm import mock_completion from litellm.caching.caching import DualCache +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 ( PresidioPerRequestConfig, OPTIONAL_PresidioPIIMasking, ) +from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import UnifiedLLMGuardrails from litellm.exceptions import GuardrailRaisedException 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, +) from litellm.exceptions import BlockedPiiEntityError @@ -696,7 +707,7 @@ 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()): assert request_data is not None presidio.add_standard_logging_guardrail_information_to_request_data( @@ -754,7 +765,7 @@ 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()): 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 +2347,7 @@ 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()): return text.replace("John Smith", "") guardrail.check_pii = mock_check_pii @@ -5070,3 +5081,196 @@ 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) + ) + ], + ) + ) + ] + ) + + +@asynccontextmanager +async def _reversible_presidio_guardrail() -> AsyncIterator[OPTIONAL_PresidioPIIMasking]: + 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("/")), + 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 + + +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": ""}') + data: Final = {"metadata": {"pii_tokens": {"": "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"}' + + +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": ""}') + ) + + await guardrail.apply_guardrail( + inputs={"texts": [], "tool_calls": [tool_call]}, + request_data={"metadata": {"pii_tokens": {"": "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": ""}' + + +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 .", "Book for Jane Doe."), ("",), id="literal-in-an-earlier-message" + ), + pytest.param(("Ask about Jane Doe.",), ("",), id="literal-beside-the-value"), + pytest.param( + ("Seen and .", "Book for Jane Doe.", "Add John Smith."), + ("", ""), + 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 diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index bd08b743c7d..f62b39a6e15 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -33083,6 +33083,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 */ From 1636c66a9ffedb95226e731dfd2271fa4d7cc24b Mon Sep 17 00:00:00 2001 From: lawrence Date: Fri, 23 Oct 2026 23:16:58 +0800 Subject: [PATCH 2/6] fix(presidio): rewrite tool call arguments one JSON value at a time An entity recognizer finds a bare value far more reliably than the same value inside JSON quoting. Presidio 2.2.364 returned no PERSON for {"text": "Jane Doe"} and only part of the name inside a longer sentence, so masking the raw arguments string let names through. Restoring by plain text replacement could also break the JSON when an original value carries a quote or a newline Arguments that parse as a JSON object or array are now rewritten one string value at a time and serialized again only when a value changed, so arguments without PII stay byte for byte the same. Anything else is still rewritten as one text --- .../guardrails/guardrail_hooks/presidio.py | 84 +++++++++++++++---- .../guardrail_hooks/test_presidio.py | 68 ++++++++++++++- 2 files changed, 134 insertions(+), 18 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index cf5a9913deb..541b86c3ce2 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -12,16 +12,25 @@ import asyncio import json import re import threading -from collections.abc import AsyncGenerator, AsyncIterable, AsyncIterator, Awaitable, Iterator, Mapping, 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, chain, count, pairwise +from itertools import accumulate, chain, count, pairwise, repeat 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 +from pydantic import ConfigDict, JsonValue, TypeAdapter, ValidationError, with_config from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm @@ -282,6 +291,37 @@ def _set_tool_call_arguments(tool_call: _ToolCall, arguments: str) -> None: pass +_JSON_VALUE: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) + + +async def _rewritten_json_strings(value: JsonValue, rewrite: Callable[[str], Awaitable[str]]) -> JsonValue: + match value: + case str(): + return await rewrite(value) + case list(): + return [await _rewritten_json_strings(item, rewrite) for item in value] + case dict(): + return {key: await _rewritten_json_strings(item, rewrite) for key, item in value.items()} + case _: + return value + + +async def _rewritten_arguments(arguments: str, rewrite: Callable[[str], 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 = _JSON_VALUE.validate_json(arguments) + except ValidationError: + return await rewrite(arguments) + if not isinstance(parsed, (dict, list)): + return await rewrite(arguments) + 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]]: @@ -1870,25 +1910,37 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): metadata: Final = (request_data.get("metadata") or {}) if request_data else {} pii_tokens: Final = metadata.get("pii_tokens", {}) - reserved: Final = accumulate(map(_literal_tokens, (source for _, source in segments)), frozenset[str].union) - rewritten: Final = ( - [self._unmask_pii_text(source, pii_tokens) for _, source 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 [ - await self.check_pii( - text=source, + async def restore(text: str) -> str: + return self._unmask_pii_text(text, pii_tokens) + + def mask_reserving(reserved_tokens: frozenset[str]) -> Callable[[str], Awaitable[str]]: + async def mask(text: str) -> 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, ) - for (_, source), reserved_tokens in zip(segments, reserved) - ] + + return mask + + rewrites: Final = ( + repeat(restore) + 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(_literal_tokens, (source for _, source in segments)), frozenset[str].union), + ) ) + rewritten: Final = [ + await (_rewritten_arguments(source, rewrite) if isinstance(slot, _ToolCallSlot) else rewrite(source)) + 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): diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_presidio.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_presidio.py index 95f1eaea0e8..b5063d99b44 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_presidio.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_presidio.py @@ -5104,9 +5104,31 @@ def _tool_call_response(arguments: str) -> ModelResponse: ) +_UNQUOTED_PERSON_NAME: Final = re.compile(r'(? 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() -> AsyncIterator[OPTIONAL_PresidioPIIMasking]: - async with TestServer(_fake_presidio_app()) as server: +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("/")), @@ -5143,6 +5165,48 @@ 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"}' +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": ""}') + data: Final = {"metadata": {"pii_tokens": {"": 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 + + 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 From 10ab70afabbc324de13a183ea7a772a43349e790 Mon Sep 17 00:00:00 2001 From: lawrence Date: Mon, 26 Oct 2026 10:07:02 +0800 Subject: [PATCH 3/6] test(presidio): cover tool call arguments that are not a JSON object Arguments that are not valid JSON are masked and restored as one text, a JSON string is rewritten as its decoded value, and a tool call without string arguments passes through untouched. The write-back helper no longer carries a branch it can never reach --- .../guardrails/guardrail_hooks/presidio.py | 15 ++++---- .../guardrail_hooks/test_presidio.py | 34 +++++++++++++++++++ 2 files changed, 40 insertions(+), 9 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index 541b86c3ce2..1b4a5419eb2 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -282,13 +282,12 @@ def _tool_call_arguments(tool_call: _ToolCall) -> str | None: def _set_tool_call_arguments(tool_call: _ToolCall, arguments: str) -> None: - match tool_call: - case {"function": function}: - function["arguments"] = arguments - case ChatCompletionMessageToolCall(function=Function() as function): - function.arguments = arguments - case _: - pass + if isinstance(tool_call, ChatCompletionMessageToolCall): + function: Final = tool_call.function + function.arguments = arguments + else: + function_chunk: Final = tool_call["function"] + function_chunk["arguments"] = arguments _JSON_VALUE: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) @@ -316,8 +315,6 @@ async def _rewritten_arguments(arguments: str, rewrite: Callable[[str], Awaitabl parsed: Final = _JSON_VALUE.validate_json(arguments) except ValidationError: return await rewrite(arguments) - if not isinstance(parsed, (dict, list)): - return await rewrite(arguments) rewritten: Final = await _rewritten_json_strings(parsed, rewrite) return arguments if rewritten == parsed else json.dumps(rewritten, ensure_ascii=False) diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_presidio.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_presidio.py index b5063d99b44..5f486ceeb83 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_presidio.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_presidio.py @@ -5207,6 +5207,40 @@ async def test_tool_call_arguments_without_pii_reach_the_model_byte_for_byte() - 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 From 4826d9f687b355df601b4cc6d24b57d6a5677614 Mon Sep 17 00:00:00 2001 From: lawrence Date: Wed, 4 Nov 2026 22:44:29 +0800 Subject: [PATCH 4/6] fix(presidio): scan structured arguments and reserve decoded tokens --- .../guardrails/guardrail_hooks/presidio.py | 64 ++++++++--- .../guardrail_hooks/test_presidio.py | 101 +++++++++++++++++- 2 files changed, 145 insertions(+), 20 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index 986f63620db..c082989f237 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -306,19 +306,47 @@ def _set_tool_call_arguments(tool_call: _ToolCall, arguments: str) -> None: _JSON_VALUE: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) -async def _rewritten_json_strings(value: JsonValue, rewrite: Callable[[str], Awaitable[str]]) -> JsonValue: +def _json_literal_tokens(value: JsonValue) -> frozenset[str]: match value: case str(): - return await rewrite(value) + 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 = _JSON_VALUE.validate_json(source) + except ValidationError: + 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 {key: await _rewritten_json_strings(item, rewrite) for key, item in value.items()} - case _: + 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], Awaitable[str]]) -> str: +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 @@ -327,7 +355,7 @@ async def _rewritten_arguments(arguments: str, rewrite: Callable[[str], Awaitabl try: parsed: Final = _JSON_VALUE.validate_json(arguments) except ValidationError: - return await rewrite(arguments) + 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) @@ -1096,7 +1124,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 @@ -1114,7 +1142,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, @@ -1127,6 +1155,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): 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 @@ -1155,7 +1184,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( @@ -1937,17 +1970,18 @@ 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) -> str: - return self._unmask_pii_text(text, 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 mask_reserving(reserved_tokens: frozenset[str]) -> Callable[[str], Awaitable[str]]: - async def mask(text: str) -> str: + 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, ) return mask @@ -1961,11 +1995,11 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): ) else map( mask_reserving, - accumulate(map(_literal_tokens, (source for _, source in segments)), frozenset[str].union), + accumulate(map(_segment_literal_tokens, segments), frozenset[str].union), ) ) rewritten: Final = [ - await (_rewritten_arguments(source, rewrite) if isinstance(slot, _ToolCallSlot) else rewrite(source)) + 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)) diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_presidio.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_presidio.py index 78ba69719ed..90e52a925bc 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_presidio.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_presidio.py @@ -707,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, reserved_tokens=frozenset()): + 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( @@ -765,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, reserved_tokens=frozenset()): + 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" @@ -2347,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, reserved_tokens=frozenset()): + 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", "") guardrail.check_pii = mock_check_pii @@ -5233,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 ) @@ -5321,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/", ) @@ -5519,6 +5535,81 @@ async def _masked_anthropic_request(guardrail: OPTIONAL_PresidioPIIMasking, mess 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 "" 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() + + +async def test_safe_tool_argument_keys_and_scalar_types_keep_original_bytes() -> None: + arguments: Final = '{ "count": 42, "fraction": 1.25, "enabled": true, "absent": null, "items": [false, 0] }' + 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 From cab0d82be1e6ceae2e79520be669619d7bf8245c Mon Sep 17 00:00:00 2001 From: lawrence Date: Thu, 5 Nov 2026 15:48:24 +0800 Subject: [PATCH 5/6] fix(presidio): reject duplicate JSON tool argument keys --- litellm/proxy/guardrails/_content_utils.py | 10 +++-- .../guardrails/guardrail_hooks/presidio.py | 21 ++++++--- .../guardrail_hooks/test_presidio.py | 45 ++++++++++++++++++- 3 files changed, 65 insertions(+), 11 deletions(-) diff --git a/litellm/proxy/guardrails/_content_utils.py b/litellm/proxy/guardrails/_content_utils.py index 51934c99db6..8034ca74f56 100644 --- a/litellm/proxy/guardrails/_content_utils.py +++ b/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) diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index c082989f237..3e00f1dd905 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -30,7 +30,7 @@ from operator import itemgetter from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, TypeAlias, cast import aiohttp -from pydantic import ConfigDict, JsonValue, TypeAdapter, ValidationError, with_config +from pydantic import ConfigDict, JsonValue, TypeAdapter, with_config from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm @@ -55,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, @@ -303,7 +304,15 @@ def _set_tool_call_arguments(tool_call: _ToolCall, arguments: str) -> None: function_chunk["arguments"] = arguments -_JSON_VALUE: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) +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]: @@ -323,8 +332,8 @@ def _segment_literal_tokens(segment: tuple[_TextSlot | _ToolCallSlot, str]) -> f if isinstance(slot, _TextSlot): return _literal_tokens(source) try: - parsed: Final = _JSON_VALUE.validate_json(source) - except ValidationError: + parsed: Final = _parsed_arguments(source) + except ValueError: return _literal_tokens(source) return _json_literal_tokens(parsed) @@ -353,8 +362,8 @@ async def _rewritten_arguments(arguments: str, rewrite: Callable[[str, bool], Aw quote or a newline has to be escaped for the arguments to stay valid JSON. """ try: - parsed: Final = _JSON_VALUE.validate_json(arguments) - except ValidationError: + 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) diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_presidio.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_presidio.py index 90e52a925bc..95bcda488ab 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_presidio.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_presidio.py @@ -5592,8 +5592,49 @@ async def test_escaped_tool_argument_tokens_never_restore_to_unrelated_pii(argum await guardrail._close_http_session() -async def test_safe_tool_argument_keys_and_scalar_types_keep_original_bytes() -> None: - arguments: Final = '{ "count": 42, "fraction": 1.25, "enabled": true, "absent": null, "items": [false, 0] }' +@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":"","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}}}, + 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( From d54353db06d619087845a4feda32b64936c722a3 Mon Sep 17 00:00:00 2001 From: lawrence Date: Thu, 5 Nov 2026 16:18:10 +0800 Subject: [PATCH 6/6] fix(presidio): restore only complete tokens in tool arguments --- .../guardrails/guardrail_hooks/presidio.py | 17 ++++++---- .../guardrail_hooks/test_presidio.py | 32 +++++++++++++++++++ 2 files changed, 43 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index 3e00f1dd905..eb9f24c6f9f 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/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" diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_presidio.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_presidio.py index 95bcda488ab..e7653397a06 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_presidio.py +++ b/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="}', + json.dumps({"url": "https://collector.invalid/?q="}', + json.dumps({"items": [{"url": "https://collector.invalid/?q="]', json.dumps([" 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}}}, + 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