From c4410080fc5312d2a8978080feda2de74ef9559a Mon Sep 17 00:00:00 2001 From: lawrence Date: Fri, 23 Oct 2026 22:41:32 +0800 Subject: [PATCH 1/3] 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/3] 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/3] 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