diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index 6436cd832c7..d636ea69f31 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, TypeVar, cast, overload, runtime_checkable @@ -519,6 +519,17 @@ def _rewritten_event(event: Mapping[str, object], rewrite_event: _SSEEventRewrit return {**event, rewrite.section: {**section, rewrite.field: rewrite.value}} +def _tool_call_text_offsets(texts_before_messages: int, extracted: Sequence[ExtractedInput]) -> tuple[int, ...]: + texts_through_message: Final = tuple( + accumulate((len(one_message.scanned) for one_message in extracted), initial=texts_before_messages) + )[1:] + return tuple( + chain.from_iterable( + repeat(offset, len(one_message.tool_calls)) for one_message, offset in zip(extracted, texts_through_message) + ) + ) + + def _tool_call_shapes(tool_calls: Sequence[object]) -> tuple[_ToolCallShape, ...]: """The guardrail-visible shape of each tool call, whether the guardrail handed back the ``ChatCompletionMessageToolCall`` objects it was given or plain dicts.""" @@ -908,7 +919,14 @@ class AnthropicMessagesHandler(BaseTranslation): # Step 2: Apply guardrail to all texts and tool calls in batch if texts_to_check or tool_calls_to_check: - inputs: Final = GenericGuardrailAPIInputs(texts=texts_to_check) + inputs: Final = ( + GenericGuardrailAPIInputs( + texts=texts_to_check, + tool_call_text_offsets=_tool_call_text_offsets(len(top_level_system_scanned), extracted), + ) + if tool_calls_to_check + else GenericGuardrailAPIInputs(texts=texts_to_check) + ) if images_to_check: inputs["images"] = images_to_check if tool_calls_to_check: diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index 5c5204ef2f4..831dae3ffd9 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 7d9e230f8ef..1edffdb7dce 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -50261,6 +50261,13 @@ "title": "Texts", "type": "array" }, + "tool_call_text_offsets": { + "items": { + "type": "integer" + }, + "title": "Tool Call Text Offsets", + "type": "array" + }, "tool_calls": { "anyOf": [ { 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/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 7e052796545..3962749f639 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 @@ -413,8 +413,7 @@ class GenericGuardrailAPI(CustomGuardrail): exc_info=error, ) # Keep flow going - treat as action=NONE (no modifications) - return_inputs: Final[GenericGuardrailAPIInputs] = {} - return_inputs.update(inputs) + return_inputs: Final[GenericGuardrailAPIInputs] = {**inputs} return return_inputs def _build_request_headers(self) -> Mapping[str, str]: diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index 458528e7e13..eb9f24c6f9f 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -12,12 +12,22 @@ import asyncio import json import re import threading -from collections.abc import AsyncGenerator, AsyncIterable, AsyncIterator, Awaitable, Iterator, Sequence +from collections.abc import ( + AsyncGenerator, + AsyncIterable, + AsyncIterator, + Awaitable, + Callable, + Iterator, + Mapping, + Sequence, +) from contextlib import asynccontextmanager from dataclasses import dataclass from datetime import datetime -from itertools import accumulate, count, pairwise -from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, cast +from itertools import accumulate, chain, count, pairwise +from operator import itemgetter +from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, TypeAlias, cast import aiohttp from pydantic import ConfigDict, JsonValue, TypeAdapter, with_config @@ -45,6 +55,7 @@ from litellm.integrations.custom_guardrail import ( ) from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.common_utils.sse_keepalive import split_complete_sse_frames +from litellm.proxy.guardrails._content_utils import DuplicateToolArgumentKeyError, parse_tool_arguments from litellm.proxy.guardrails.anthropic_sse import ( anthropic_sse_chunks_from_response, assemble_anthropic_sse_stream, @@ -59,11 +70,12 @@ from litellm.types.guardrails import ( PiiEntityType, PresidioPerRequestConfig, ) +from litellm.types.llms.openai import ChatCompletionToolCallChunk from litellm.types.proxy.guardrails.guardrail_hooks.presidio import ( PresidioAnalyzeRequest, PresidioAnalyzeResponseItem, ) -from litellm.types.utils import GuardrailStatus, Message, StreamingChoices +from litellm.types.utils import ChatCompletionMessageToolCall, Function, GuardrailStatus, Message, StreamingChoices from litellm.utils import ( EmbeddingResponse, ImageResponse, @@ -249,6 +261,138 @@ def _resolve_overlapping_spans(analyze_results: Sequence[PresidioAnalyzeResponse return tuple(_collapse_run(detections[first:last]) for first, last in pairwise(boundaries) if first < last) +_LITERAL_TOKEN: Final = re.compile(r"<[^<>\s]+_\d+>") + +_ToolCall: TypeAlias = ChatCompletionToolCallChunk | ChatCompletionMessageToolCall + + +@dataclass(frozen=True, slots=True) +class _TextSlot: + index: int + + +@dataclass(frozen=True, slots=True) +class _ToolCallSlot: + index: int + + +def _literal_tokens(text: str) -> frozenset[str]: + return frozenset(text[match.start() : match.end()] for match in _LITERAL_TOKEN.finditer(text)) + + +def _free_token(entity_type: str, numbers: Iterator[int], held: Mapping[str, str], reserved: frozenset[str]) -> str: + candidates: Final = (f"<{entity_type}_{number}>" for number in numbers) + return next(token for token in candidates if token not in held and token not in reserved) + + +def _tool_call_arguments(tool_call: _ToolCall) -> str | None: + match tool_call: + case {"function": {"arguments": str() as arguments}}: + return arguments + case ChatCompletionMessageToolCall(function=Function(arguments=str() as arguments)): + return arguments + case _: + return None + + +def _set_tool_call_arguments(tool_call: _ToolCall, arguments: str) -> None: + if isinstance(tool_call, ChatCompletionMessageToolCall): + function: Final = tool_call.function + function.arguments = arguments + else: + function_chunk: Final = tool_call["function"] + function_chunk["arguments"] = arguments + + +def _parsed_arguments(arguments: str) -> JsonValue: + try: + return parse_tool_arguments(arguments) + except DuplicateToolArgumentKeyError as error: + raise GuardrailRaisedException( + message="Presidio cannot safely scan tool arguments with duplicate JSON object keys", + should_wrap_with_default_message=False, + blocked_content=True, + ) from error + + +def _json_literal_tokens(value: JsonValue) -> frozenset[str]: + match value: + case str(): + return _literal_tokens(value) + case list(): + return frozenset[str]().union(*map(_json_literal_tokens, value)) + case dict(): + return frozenset[str]().union(*map(_literal_tokens, value), *map(_json_literal_tokens, value.values())) + case _: + return frozenset() + + +def _segment_literal_tokens(segment: tuple[_TextSlot | _ToolCallSlot, str]) -> frozenset[str]: + slot, source = segment + if isinstance(slot, _TextSlot): + return _literal_tokens(source) + try: + parsed: Final = _parsed_arguments(source) + except ValueError: + return _literal_tokens(source) + return _json_literal_tokens(parsed) + + +async def _rewritten_json_strings(value: JsonValue, rewrite: Callable[[str, bool], Awaitable[str]]) -> JsonValue: + match value: + case str(): + return await rewrite(value, True) + case list(): + return [await _rewritten_json_strings(item, rewrite) for item in value] + case dict(): + return { + await rewrite(key, False): await _rewritten_json_strings(item, rewrite) for key, item in value.items() + } + case bool() | None: + return value + case int() | float(): + await rewrite(str(value), False) + return value + + +async def _rewritten_arguments(arguments: str, rewrite: Callable[[str, bool], Awaitable[str]]) -> str: + """ + Rewrite each string of a tool call's JSON arguments on its own. An entity recognizer + finds a bare value that it misses inside JSON quoting, and a restored value with a + quote or a newline has to be escaped for the arguments to stay valid JSON. + """ + try: + parsed: Final = _parsed_arguments(arguments) + except ValueError: + return await rewrite(arguments, True) + rewritten: Final = await _rewritten_json_strings(parsed, rewrite) + return arguments if rewritten == parsed else json.dumps(rewritten, ensure_ascii=False) + + +def _conversation_segments( + texts: Sequence[str], arguments: Sequence[str | None], tool_call_text_offsets: Sequence[int] +) -> Iterator[tuple[_TextSlot | _ToolCallSlot, str]]: + """ + Every text and tool call argument string in the order the conversation holds them, + so masking an earlier message never depends on the turns appended after it. + """ + offsets: Final = ( + tool_call_text_offsets if len(tool_call_text_offsets) == len(arguments) else (len(texts),) * len(arguments) + ) + keyed: Final = sorted( + chain( + (((index, 1), _TextSlot(index), text) for index, text in enumerate(texts)), + ( + ((offset, 0), _ToolCallSlot(index), source) + for index, (offset, source) in enumerate(zip(offsets, arguments)) + if source is not None + ), + ), + key=itemgetter(0), + ) + return ((slot, source) for _, slot, source in keyed) + + class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): user_api_key_cache = None ad_hoc_recognizers: list[str] | None = None @@ -872,6 +1016,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): analyze_results: Any, request_data: dict | None, masked_entity_count: dict[str, int], + reserved_tokens: frozenset[str] = frozenset(), ) -> str: # output_parse_pii is True — we need sequentially numbered # tokens and a pii_tokens mapping for later unmasking. @@ -893,8 +1038,8 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): runs: Final = _resolve_overlapping_spans(analyze_results) # Every message and content part numbers its runs from 1 into this one map, so a - # token another part already holds for a different value takes a spare number. - # No token in the map is numbered above its size, so spares start above both. + # token another part already holds for a different value, or one the request already + # carries as literal text, takes a spare number that is neither held nor reserved. spare_numbers: Final = count(max(len(pii_tokens), len(runs)) + 1) for seq, run in reversed(tuple(enumerate(runs, start=1))): @@ -902,8 +1047,8 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): numbered = f"<{run.entity_type}_{seq}>" replacement = ( numbered - if pii_tokens.get(numbered, original) == original - else f"<{run.entity_type}_{next(spare_numbers)}>" + if numbered not in reserved_tokens and pii_tokens.get(numbered, original) == original + else _free_token(run.entity_type, spare_numbers, pii_tokens, reserved_tokens) ) pii_tokens[replacement] = original new_text = new_text[: run.start] + replacement + new_text[run.end :] @@ -917,6 +1062,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): output_parse_pii: bool, masked_entity_count: dict[str, int], request_data: dict | None = None, + reserved_tokens: frozenset[str] = frozenset(), ) -> str: """ Send analysis results to the Presidio anonymizer endpoint to get redacted text @@ -936,7 +1082,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): return self._finalize_presidio_anonymize_simple(redacted_text, masked_entity_count) return self._finalize_presidio_anonymize_numbered_tokens( - text, analyze_results, request_data, masked_entity_count + text, analyze_results, request_data, masked_entity_count, reserved_tokens ) except Exception as e: # Sanitize exception to avoid leaking the original text (which may @@ -987,7 +1133,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): return filtered_results def raise_exception_if_blocked_entities_detected( - self, analyze_results: list[PresidioAnalyzeResponseItem] | _PresidioAnonymizeResponse + self, analyze_results: list[PresidioAnalyzeResponseItem] | _PresidioAnonymizeResponse, block_all: bool = False ) -> None: """ Raise an exception if blocked entities are detected @@ -1005,7 +1151,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): if entity_type: # Check if entity_type is in config (supports both enum and string) - if entity_type in self.pii_entities_config and self.pii_entities_config[entity_type] == PiiAction.BLOCK: + if block_all or self.pii_entities_config.get(entity_type) == PiiAction.BLOCK: raise BlockedPiiEntityError( entity_type=entity_type, guardrail_name=self.guardrail_name, @@ -1017,6 +1163,8 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): output_parse_pii: bool, presidio_config: PresidioPerRequestConfig | None, request_data: dict, + reserved_tokens: frozenset[str] = frozenset(), + allow_masking: bool = True, ) -> str: """ Calls Presidio Analyze + Anonymize endpoints for PII Analysis + Masking @@ -1045,7 +1193,11 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): #################################################### # Blocked Entities check #################################################### - self.raise_exception_if_blocked_entities_detected(analyze_results=analyze_results) + self.raise_exception_if_blocked_entities_detected( + analyze_results=analyze_results, block_all=not allow_masking + ) + if not allow_masking: + return text # Then anonymize the text using the analysis results anonymized_text: Final = await self.anonymize_text( @@ -1054,6 +1206,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): output_parse_pii=output_parse_pii, masked_entity_count=masked_entity_count, request_data=request_data, + reserved_tokens=reserved_tokens, ) return anonymized_text return redacted_text["text"] @@ -1314,7 +1467,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): return response @staticmethod - def _unmask_pii_text(text: str, pii_tokens: dict[str, str]) -> str: + def _unmask_pii_text(text: str, pii_tokens: dict[str, str], *, allow_partial_tokens: bool = True) -> str: """ Replace PII tokens in *text* with their original values. @@ -1328,7 +1481,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): for token, original_text in pii_tokens.items(): if token in text: text = text.replace(token, original_text) - else: + elif allow_partial_tokens: # FALLBACK: Handle truncated tokens (token cut off by max_tokens) # Only check at the very end of the text. min_overlap = min(20, len(token) // 2) @@ -1816,27 +1969,58 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): 2. When Testing the guardrail with some text, this function will be called with the input text and returns a text after applying the guardrail """ texts: Final = inputs.get("texts", []) + tool_calls: Final = inputs.get("tool_calls") or () + segments: Final = tuple( + _conversation_segments( + texts, tuple(map(_tool_call_arguments, tool_calls)), inputs.get("tool_call_text_offsets", ()) + ) + ) metadata: Final = (request_data.get("metadata") or {}) if request_data else {} pii_tokens: Final = metadata.get("pii_tokens", {}) - new_texts: Final = [] - if input_type == "response" and ( - self._callback_role == "restore" - or (self._callback_role is None and not self.apply_to_output and pii_tokens) - ): - for text in texts: - new_texts.append(self._unmask_pii_text(text, pii_tokens)) - else: - for text in texts: - modified_text = await self.check_pii( + def restore_for_slot(slot: _TextSlot | _ToolCallSlot) -> Callable[[str, bool], Awaitable[str]]: + async def restore(text: str, allow_masking: bool) -> str: + if not allow_masking: + return text + return self._unmask_pii_text(text, pii_tokens, allow_partial_tokens=isinstance(slot, _TextSlot)) + + return restore + + def mask_reserving(reserved_tokens: frozenset[str]) -> Callable[[str, bool], Awaitable[str]]: + async def mask(text: str, allow_masking: bool) -> str: + return await self.check_pii( text=text, output_parse_pii=self.output_parse_pii, presidio_config=None, request_data=request_data or {}, + reserved_tokens=reserved_tokens, + allow_masking=allow_masking, ) - new_texts.append(modified_text) - inputs["texts"] = new_texts + + return mask + + rewrites: Final = ( + map(restore_for_slot, (slot for slot, _ in segments)) + if input_type == "response" + and ( + self._callback_role == "restore" + or (self._callback_role is None and not self.apply_to_output and pii_tokens) + ) + else map( + mask_reserving, + accumulate(map(_segment_literal_tokens, segments), frozenset[str].union), + ) + ) + rewritten: Final = [ + await (_rewritten_arguments(source, rewrite) if isinstance(slot, _ToolCallSlot) else rewrite(source, True)) + for (slot, source), rewrite in zip(segments, rewrites) + ] + by_slot: Final = dict(zip((slot for slot, _ in segments), rewritten)) + inputs["texts"] = [by_slot[_TextSlot(index)] for index in range(len(texts))] + for index, tool_call in enumerate(tool_calls): + if _ToolCallSlot(index) in by_slot: + _set_tool_call_arguments(tool_call, by_slot[_ToolCallSlot(index)]) return inputs def update_in_memory_litellm_params(self, litellm_params: LitellmParams) -> None: diff --git a/litellm/proxy/guardrails/guardrail_hooks/straiker/straiker.py b/litellm/proxy/guardrails/guardrail_hooks/straiker/straiker.py index 589b865618b..3fee83e49b3 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 ba29a67a36c..179233fc699 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -5361,6 +5361,7 @@ class GenericGuardrailAPIInputs(TypedDict, total=False): images: list[str] # extracted images from the LLM response - for image guardrails tools: list[ChatCompletionToolParam] # tools sent to the LLM tool_calls: list[ChatCompletionToolCallChunk] | list[ChatCompletionMessageToolCall] # tool calls sent from the LLM + tool_call_text_offsets: ReadOnly[tuple[int, ...]] # per tool call, how many texts precede it in the conversation structured_messages: list[ AllMessageValues ] # structured messages sent to the LLM - indicates if text is from system or user diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_presidio.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_presidio.py index 563b2a0f2b0..e7653397a06 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_presidio.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_presidio.py @@ -26,14 +26,25 @@ import litellm from litellm import mock_completion from litellm.caching.caching import DualCache from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException +from litellm.llms.anthropic.chat.guardrail_translation.handler import AnthropicMessagesHandler +from litellm.llms.openai.chat.guardrail_translation.handler import OpenAIChatCompletionsHandler from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.guardrails.guardrail_hooks.presidio import ( OPTIONAL_PresidioPIIMasking, PresidioPerRequestConfig, ) +from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import UnifiedLLMGuardrails from litellm.types.guardrails import LitellmParams, PiiAction, PiiEntityType from litellm.types.proxy.guardrails.guardrail_hooks.presidio import PresidioAnalyzeRequest, PresidioAnalyzeResponseItem -from litellm.types.utils import Choices, Delta, Message, ModelResponse, StreamingChoices +from litellm.types.utils import ( + ChatCompletionMessageToolCall, + Choices, + Delta, + Function, + Message, + ModelResponse, + StreamingChoices, +) def _make_mock_session_iterator(json_response, status=200, content_type="application/json", text_response=""): @@ -696,7 +707,9 @@ async def test_presidio_sets_guardrail_information_in_request_data(): "metadata": {}, } - async def mock_check_pii(text, output_parse_pii, presidio_config, request_data): + async def mock_check_pii( + text, output_parse_pii, presidio_config, request_data, reserved_tokens=frozenset(), *, allow_masking: bool = True + ): assert request_data is not None presidio.add_standard_logging_guardrail_information_to_request_data( @@ -754,7 +767,9 @@ async def test_request_data_flows_to_apply_guardrail(): "metadata": {}, } - async def mock_check_pii(text, output_parse_pii, presidio_config, request_data): + async def mock_check_pii( + text, output_parse_pii, presidio_config, request_data, reserved_tokens=frozenset(), *, allow_masking: bool = True + ): assert request_data is not None, "request_data should be passed to check_pii" assert "metadata" in request_data, "request_data should have metadata" @@ -2336,7 +2351,9 @@ async def test_standalone_scans_without_restoration_tokens(input_type: Literal[" mock_testing=True, ) - async def mock_check_pii(text, output_parse_pii, presidio_config, request_data): + async def mock_check_pii( + text, output_parse_pii, presidio_config, request_data, reserved_tokens=frozenset(), *, allow_masking: bool = True + ): return text.replace("John Smith", "") guardrail.check_pii = mock_check_pii @@ -5222,7 +5239,11 @@ def _card_and_email_spans(text: str) -> tuple[PresidioAnalyzeResponseItem, ...]: score=1.0, analysis_explanation=None, ) - for entity_type, value in (("CREDIT_CARD", _CARD_NUMBER), ("EMAIL_ADDRESS", _EMAIL)) + for entity_type, value in ( + ("CREDIT_CARD", _CARD_NUMBER), + ("CREDIT_CARD", _CARD_NUMBER.replace("-", "")), + ("EMAIL_ADDRESS", _EMAIL), + ) if value in text ) @@ -5310,9 +5331,15 @@ def _drain(analyzed: asyncio.Queue[PresidioAnalyzeRequest]) -> tuple[PresidioAna def _guardrail_with_in_memory_presidio( analyzed: asyncio.Queue[PresidioAnalyzeRequest], + *, + output_parse_pii: bool = False, + action: PiiAction | None = None, ) -> OPTIONAL_PresidioPIIMasking: guardrail: Final = OPTIONAL_PresidioPIIMasking( - pii_entities_config=_BLOCK_CARD_MASK_EMAIL, + output_parse_pii=output_parse_pii, + pii_entities_config=( + _BLOCK_CARD_MASK_EMAIL if action is None else {entity: action for entity in _BLOCK_CARD_MASK_EMAIL} + ), presidio_analyzer_api_base="http://presidio-analyzer.test/", presidio_anonymizer_api_base="http://presidio-anonymizer.test/", ) @@ -5438,3 +5465,442 @@ def test_get_session_iterator_reuses_one_session_per_background_loop( assert len(sessions) == 10 assert all(session is sessions[0] for session in sessions) assert presidio_guardrail._http_session is None + +def _function_tool_call(call_id: str, arguments: str) -> dict: + return {"id": call_id, "type": "function", "function": {"name": "lookup", "arguments": arguments}} + + +def _tool_call_response(arguments: str) -> ModelResponse: + return ModelResponse( + choices=[ + Choices( + message=Message( + role="assistant", + content=None, + tool_calls=[ + ChatCompletionMessageToolCall( + id="call_1", type="function", function=Function(name="lookup", arguments=arguments) + ) + ], + ) + ) + ] + ) + + +_UNQUOTED_PERSON_NAME: Final = re.compile(r'(? web.Response: + payload = await request.json() + return web.json_response( + [ + {"entity_type": "PERSON", "start": match.start(), "end": match.end(), "score": 0.85} + for match in _UNQUOTED_PERSON_NAME.finditer(payload["text"]) + ] + ) + + +def _json_blind_presidio_app() -> web.Application: + app = web.Application() + app.router.add_post("/analyze", _analyze_names_outside_json_quotes) + app.router.add_post("/anonymize", _fake_anonymize) + return app + + +@asynccontextmanager +async def _reversible_presidio_guardrail( + presidio_app: web.Application | None = None, +) -> AsyncIterator[OPTIONAL_PresidioPIIMasking]: + async with TestServer(presidio_app or _fake_presidio_app()) as server: + guardrail: Final = OPTIONAL_PresidioPIIMasking( + presidio_analyzer_api_base=str(server.make_url("/")), + presidio_anonymizer_api_base=str(server.make_url("/")), + output_parse_pii=True, + pii_entities_config={PiiEntityType.PERSON: PiiAction.MASK}, + ) + yield guardrail + await guardrail._close_http_session() + + +async def _masked_chat_request(guardrail: OPTIONAL_PresidioPIIMasking, messages: Sequence[dict]) -> dict: + data: Final = {"model": "claude-fable-5-1", "messages": copy.deepcopy(list(messages))} + await OpenAIChatCompletionsHandler().process_input_messages(data=data, guardrail_to_apply=guardrail) + return data + + +async def _masked_anthropic_request(guardrail: OPTIONAL_PresidioPIIMasking, messages: Sequence[dict]) -> dict: + data: Final = {"model": "claude-fable-5-1", "max_tokens": 64, "messages": copy.deepcopy(list(messages))} + await AnthropicMessagesHandler().process_input_messages(data=data, guardrail_to_apply=guardrail) + return data + + +@pytest.mark.parametrize("action", [PiiAction.MASK, PiiAction.BLOCK]) +@pytest.mark.parametrize( + "arguments", + [ + json.dumps({_EMAIL: {"safe": True}}), + json.dumps({"nested": [{"card": int(_CARD_NUMBER.replace("-", ""))}]}), + _CARD_NUMBER.replace("-", ""), + _CARD_NUMBER.replace("-", "") + ".0", + ], +) +async def test_tool_argument_keys_and_numbers_cannot_bypass_pii_policy(arguments: str, action: PiiAction) -> None: + analyzed: Final[asyncio.Queue[PresidioAnalyzeRequest]] = asyncio.Queue() + guardrail: Final = _guardrail_with_in_memory_presidio(analyzed, output_parse_pii=True, action=action) + tool_call: Final = ChatCompletionMessageToolCall( + id="call_1", type="function", function=Function(name="lookup", arguments=arguments) + ) + try: + with pytest.raises(BlockedPiiEntityError) as excinfo: + await guardrail.apply_guardrail( + inputs={"texts": [], "tool_calls": [tool_call]}, + request_data={"metadata": {}}, + input_type="request", + ) + finally: + await guardrail._close_http_session() + assert excinfo.value.status_code == 400 + assert excinfo.value.entity_type == ("EMAIL_ADDRESS" if _EMAIL in arguments else "CREDIT_CARD") + + +@pytest.mark.parametrize( + "arguments", + [ + r'{"guest":"test@example.com","nested":[{"url":"https://example.invalid/?q=\u003cEMAIL_ADDRESS_1\u003e"}]}', + r'{"guest":"test@example.com","\u003cEMAIL_ADDRESS_1\u003e":null}', + ], +) +async def test_escaped_tool_argument_tokens_never_restore_to_unrelated_pii(arguments: str) -> None: + analyzed: Final[asyncio.Queue[PresidioAnalyzeRequest]] = asyncio.Queue() + guardrail: Final = _guardrail_with_in_memory_presidio(analyzed, output_parse_pii=True) + request_data: Final[dict[str, dict[str, dict[str, str]]]] = {"metadata": {}} + tool_call: Final = ChatCompletionMessageToolCall( + id="call_1", type="function", function=Function(name="lookup", arguments=arguments) + ) + try: + masked: Final = await guardrail.apply_guardrail( + inputs={"texts": [], "tool_calls": [tool_call]}, + request_data=request_data, + input_type="request", + ) + assert _EMAIL not in tool_call.function.arguments + assert "" not in request_data["metadata"]["pii_tokens"] + await guardrail.apply_guardrail(inputs=masked, request_data=request_data, input_type="response") + assert json.loads(tool_call.function.arguments) == json.loads(arguments) + finally: + await guardrail._close_http_session() + + +@pytest.mark.parametrize( + ("input_type", "action"), + [("request", PiiAction.MASK), ("request", PiiAction.BLOCK), ("response", PiiAction.MASK)], +) +@pytest.mark.parametrize( + "arguments", + [ + '{"guest":"test@example.com","guest":"safe"}', + '{"items":[{"guest":"test@example.com","guest":"safe"}]}', + r'{"guest":"test@example.com","\u0067uest":"safe"}', + '{"guest":"","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( + id="call_1", type="function", function=Function(name="lookup", arguments=arguments) + ) + try: + await guardrail.apply_guardrail( + inputs={"texts": [], "tool_calls": [tool_call]}, + request_data={"metadata": {}}, + input_type="request", + ) + assert tool_call.function.arguments == arguments + finally: + await guardrail._close_http_session() + + +async def test_restore_writes_the_original_values_into_non_streaming_tool_call_arguments() -> None: + guardrail: Final = OPTIONAL_PresidioPIIMasking( + event_hook="pre_call", default_on=True, output_parse_pii=True, mock_testing=True + ) + response: Final = _tool_call_response('{"guest": ""}') + 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"}' + + +@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 + ) + 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 + + +@pytest.mark.parametrize( + "arguments", + [ + pytest.param("call Jane Doe back", id="not-json"), + pytest.param('"Jane Doe"', id="json-string"), + ], +) +async def test_tool_call_arguments_that_are_not_a_json_object_are_still_masked_and_restored(arguments: str) -> None: + async with _reversible_presidio_guardrail() as guardrail: + data: Final = await _masked_chat_request( + guardrail, + [{"role": "assistant", "content": None, "tool_calls": [_function_tool_call("call_1", arguments)]}], + ) + + masked: Final = data["messages"][0]["tool_calls"][0]["function"]["arguments"] + assert "Jane Doe" not in masked, masked + assert OPTIONAL_PresidioPIIMasking._unmask_pii_text(masked, data["metadata"]["pii_tokens"]) == arguments, masked + + +async def test_a_tool_call_without_string_arguments_passes_through_beside_masked_text() -> None: + tool_call: Final = {"id": "call_1", "type": "function", "function": {"name": "lookup"}} + async with _reversible_presidio_guardrail() as guardrail: + data: Final = await _masked_chat_request( + guardrail, + [ + {"role": "user", "content": "Book for Jane Doe."}, + {"role": "assistant", "content": None, "tool_calls": [tool_call]}, + ], + ) + + assert "Jane Doe" not in data["messages"][0]["content"], data["messages"][0] + assert data["messages"][1]["tool_calls"] == [tool_call] + + +async def test_apply_guardrail_restores_the_arguments_of_tool_call_objects() -> None: + guardrail: Final = OPTIONAL_PresidioPIIMasking( + event_hook="pre_call", default_on=True, output_parse_pii=True, mock_testing=True + ) + tool_call: Final = ChatCompletionMessageToolCall( + id="call_1", type="function", function=Function(name="lookup", arguments='{"guest": ""}') + ) + + 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 d14e490f606..c8c154c4003 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -34185,6 +34185,8 @@ export interface components { structured_messages?: (components["schemas"]["ChatCompletionUserMessage"] | components["schemas"]["ChatCompletionAssistantMessage"] | components["schemas"]["ChatCompletionToolMessage"] | components["schemas"]["ChatCompletionSystemMessage"] | components["schemas"]["ChatCompletionFunctionMessage"] | components["schemas"]["ChatCompletionDeveloperMessage"])[]; /** Texts */ texts?: string[]; + /** Tool Call Text Offsets */ + tool_call_text_offsets?: number[]; /** Tool Calls */ tool_calls?: components["schemas"]["ChatCompletionToolCallChunk"][] | components["schemas"]["ChatCompletionMessageToolCall"][]; /** Tools */