diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index 7e3fda06b1a..b495d606e73 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.""" @@ -888,7 +899,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 5cc1c53d396..513b0231ae7 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -49772,6 +49772,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 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..986f63620db 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -12,15 +12,25 @@ 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, 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 @@ -59,11 +69,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 +260,102 @@ 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 + + +_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) + 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 +979,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 +1001,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 +1010,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 +1025,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 +1045,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 @@ -1017,6 +1126,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 @@ -1054,6 +1164,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"] @@ -1816,27 +1927,52 @@ 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( + 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, ) - new_texts.append(modified_text) - inputs["texts"] = new_texts + + 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): + 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 ee80b63f93b..6d9f6e723bd 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -5353,6 +5353,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..78ba69719ed 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,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 @@ -5438,3 +5449,294 @@ 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 + + +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_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 e2e1210bc2a..22195a55e94 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -33613,6 +33613,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 */