From b8e7f13aeb0f71e6afa5ace7d2b72c1e4e0984f2 Mon Sep 17 00:00:00 2001 From: jesus Date: Tue, 3 Nov 2026 22:52:52 +0000 Subject: [PATCH] fix(guardrails): write rewritten tool_calls from the Generic Guardrail API back to the client The Generic Guardrail API dropped the tool_calls a post_call guardrail returned, and the Anthropic /v1/messages output handler never wrote returned tool calls back onto tool_use blocks, so a guardrail that masked tool-call arguments left the original arguments in the response on /v1/chat/completions, /v1/responses and /v1/messages. Absorbs the fix from #45161 Co-authored-by: Ivana Dunisijevic --- .../chat/guardrail_translation/handler.py | 189 +++++++++++++- .../chat/guardrail_translation/handler.py | 163 +++++++++++-- .../generic_guardrail_api.py | 2 + .../guardrail_hooks/generic_guardrail_api.py | 47 +++- .../test_anthropic_guardrail_handler.py | 230 ++++++++++++++++++ .../test_openai_guardrail_handler.py | 145 +++++++++++ .../test_generic_guardrail_api.py | 60 +++++ 7 files changed, 802 insertions(+), 34 deletions(-) diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index 7e3fda06b1a..8ba69a6132e 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -155,6 +155,12 @@ class _TextAttrBlock(Protocol): text: str +@runtime_checkable +class _ToolUseAttrBlock(Protocol): + name: str + input: Mapping[str, object] + + @runtime_checkable class _WritableMessage(Protocol): @overload @@ -166,6 +172,18 @@ class _WritableMessage(Protocol): def __setitem__(self, key: str, value: object, /) -> None: ... +_MESSAGE_ROWS: Final = TypeAdapter(tuple[object, ...]) + + +def _message_rows(data: Mapping[str, object]) -> Sequence[_WritableMessage]: + """The request's own message rows, so a write-back lands on the request itself.""" + try: + rows: Final = _MESSAGE_ROWS.validate_python(data.get("messages")) + except ValidationError: + return () + return cast(Sequence[_WritableMessage], rows) # cast-ok: the rows are the request's message dicts + + def _as_writable(value: _WritableMessage) -> _WritableMessage: return value @@ -435,8 +453,45 @@ _RELEASED_TOOL_USE_STOP: Final = ( def _rewritten_tool_use_input(arguments: str) -> Mapping[str, object] | None: + """OpenAI-shaped ``arguments`` for a call with no parameters are ``""``, which becomes the empty + ``tool_use`` input, as ``input: {}`` is what ``arguments: ""`` was translated from.""" try: - return _TOOL_USE_INPUT_ADAPTER.validate_json(arguments) + return _TOOL_USE_INPUT_ADAPTER.validate_json("{}" if arguments == "" else arguments) + except ValidationError: + return None + + +_CONTENT_BLOCKS: Final = TypeAdapter(tuple[object, ...]) +_MAPPING_BLOCK: Final = TypeAdapter(dict[str, object]) + + +def _tool_use_targets(messages: Sequence[_WritableMessage]) -> tuple[ToolUseInputTarget, ...]: + """Where the client's tool_use blocks sit now, since a structured-message write-back can merge + text blocks and move a tool_use away from the position it was scanned at.""" + return tuple( + chain.from_iterable(_row_tool_use_targets(msg_idx, message) for msg_idx, message in enumerate(messages)) + ) + + +def _row_tool_use_targets(msg_idx: int, message: _WritableMessage) -> tuple[ToolUseInputTarget, ...]: + return tuple( + ToolUseInputTarget(msg_idx, content_idx) + for content_idx, block in enumerate(_content_blocks(message)) + if block is not None and _is_client_tool_use(block) + ) + + +def _content_blocks(message: _WritableMessage) -> tuple[Mapping[str, object] | None, ...]: + try: + blocks: Final = _CONTENT_BLOCKS.validate_python(message.get("content")) + except ValidationError: + return () + return tuple(_mapping_block(block) for block in blocks) + + +def _mapping_block(block: object) -> Mapping[str, object] | None: + try: + return _MAPPING_BLOCK.validate_python(block) except ValidationError: return None @@ -453,6 +508,22 @@ def _write_back_tool_use( block["name"] = shape.name +def _write_back_tool_use_block( + block: object, before: _ToolCallShape, after: _ToolCallShape, rewritten_input: Mapping[str, object] +) -> None: + renamed: Final = after.name is not None and after.name != before.name + if isinstance(block, dict): + block["input"] = rewritten_input # rebind-ok: the tool_use block is the client-visible object being rewritten + if renamed: + block["name"] = after.name # rebind-ok: same in-place rewrite as the input above + return + if not isinstance(block, _ToolUseAttrBlock): + return + block.input = rewritten_input # rebind-ok: same in-place rewrite for an object-shaped tool_use block + if renamed and after.name is not None: + block.name = after.name # rebind-ok: same in-place rewrite as the input above + + def _streamed_query_strings( tool_call_blocks: Sequence[tuple[int, Mapping[str, object]]], tool_calls: Sequence[_ToolCallShape] ) -> tuple[_WebSearchString, ...]: @@ -532,6 +603,27 @@ def _tool_call_shapes(tool_calls: Sequence[object]) -> tuple[_ToolCallShape, ... ) +def _post_guardrail_tool_call_shapes( + returned_tool_calls: Sequence[object] | None, + pre_guardrail_tool_calls: tuple[_ToolCallShape, ...], + scanned_tool_calls: Sequence[object], + guardrail_name: str | None, +) -> tuple[_ToolCallShape, ...]: + """A returned list of another length cannot be matched to the scanned calls by position, so + the calls the guardrail was handed stand, as on the chat and responses endpoints.""" + if returned_tool_calls is None: + return _tool_call_shapes(scanned_tool_calls) + if len(returned_tool_calls) == len(pre_guardrail_tool_calls): + return _tool_call_shapes(returned_tool_calls) + verbose_proxy_logger.warning( + "Anthropic Messages: guardrail %s returned %d tool calls for the %d scanned, leaving the tool_use blocks unchanged", + guardrail_name, + len(returned_tool_calls), + len(pre_guardrail_tool_calls), + ) + return _tool_call_shapes(scanned_tool_calls) + + @with_config(ConfigDict(extra="ignore", strict=True)) class _SSETextDelta(TypedDict, total=False): type: ReadOnly[object] @@ -943,6 +1035,7 @@ class AnthropicMessagesHandler(BaseTranslation): guardrailed_structured_messages is not None and guardrailed_structured_messages is not original_structured_messages ): + scanned_rows: Final = _message_rows(cast(Mapping[str, object], data)) # cast-ok: Anthropic request body self._write_back_structured_messages( data, guardrailed_structured_messages @@ -956,6 +1049,25 @@ class AnthropicMessagesHandler(BaseTranslation): preserve_system_messages=has_midturn_system_message, original_tool_names=original_tool_names.used, ) + rewritten_messages: Final = _message_rows( + cast(Mapping[str, object], data) # cast-ok: Anthropic request body + ) + if len(rewritten_messages) == len(scanned_rows): + self._apply_guardrail_tool_calls_to_input( + messages=rewritten_messages, + scanned_tool_calls=scanned_tool_calls, + pre_guardrail_tool_calls=pre_guardrail_tool_calls, + returned_tool_calls=guardrailed_inputs.get("tool_calls"), + guardrail_name=guardrail_to_apply.guardrail_name, + ) + else: + verbose_proxy_logger.warning( + "Anthropic Messages: guardrail %s rewrote the conversation to %d rows from %d, " + "so its tool_calls rewrite cannot be placed and is left unapplied", + guardrail_to_apply.guardrail_name, + len(rewritten_messages), + len(scanned_rows), + ) else: if guardrailed_texts and len(guardrailed_texts) != len(scanned): raise unappliable_request_rewrite(guardrail_to_apply.guardrail_name) @@ -1402,23 +1514,36 @@ class AnthropicMessagesHandler(BaseTranslation): returned_tool_calls: Sequence[object] | None, guardrail_name: str | None, ) -> None: - post_guardrail_tool_calls: Final = _tool_call_shapes( - returned_tool_calls - if returned_tool_calls is not None and len(returned_tool_calls) == len(pre_guardrail_tool_calls) - else tuple(item.tool_call for item in scanned_tool_calls) + post_guardrail_tool_calls: Final = _post_guardrail_tool_call_shapes( + returned_tool_calls, + pre_guardrail_tool_calls, + tuple(item.tool_call for item in scanned_tool_calls), + guardrail_name, ) + targets: Final = _tool_use_targets(messages) + if len(targets) != len(scanned_tool_calls): + verbose_proxy_logger.warning( + "Anthropic Messages: guardrail %s left %d tool_use blocks where %d were scanned, " + "so its tool_calls rewrite cannot be placed and is left unapplied", + guardrail_name, + len(targets), + len(scanned_tool_calls), + ) + return rewritten: Final = tuple( - (item, after, _rewritten_tool_use_input(after.arguments)) - for item, before, after in zip(scanned_tool_calls, pre_guardrail_tool_calls, post_guardrail_tool_calls) + (target, after, _rewritten_tool_use_input(after.arguments)) + for target, before, after in zip(targets, pre_guardrail_tool_calls, post_guardrail_tool_calls) if before != after ) applicable: Final = tuple( - (item, after, rewritten_input) for item, after, rewritten_input in rewritten if rewritten_input is not None + (target, after, rewritten_input) + for target, after, rewritten_input in rewritten + if rewritten_input is not None ) if len(applicable) != len(rewritten): raise unappliable_request_rewrite(guardrail_name) - for item, after, rewritten_input in applicable: - _write_back_tool_use(messages[item.target.msg_idx], item.target, after, rewritten_input) + for target, after, rewritten_input in applicable: + _write_back_tool_use(messages[target.msg_idx], target, after, rewritten_input) async def process_output_response( self, @@ -1488,6 +1613,7 @@ class AnthropicMessagesHandler(BaseTranslation): tool_calls_to_check, response, ) + pre_guardrail_tool_calls: Final = _tool_call_shapes(tool_calls_to_check) guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail( inputs=inputs, @@ -1509,6 +1635,13 @@ class AnthropicMessagesHandler(BaseTranslation): self._write_back_web_search_strings( response_content, _web_search_rewrites_by_block(web_search_strings, web_search_rewrites) ) + self._apply_guardrail_tool_calls_to_output( + response_content, + tool_calls_to_check, + pre_guardrail_tool_calls, + guardrailed_inputs.get("tool_calls"), + guardrail_to_apply.guardrail_name, + ) verbose_proxy_logger.debug("Anthropic Messages: Processed output response: %s", response) @@ -1682,6 +1815,42 @@ class AnthropicMessagesHandler(BaseTranslation): for block_idx, rewrites in rewrites_by_block.items(): _write_back_web_search_block(_content_block_at(response_content, block_idx), rewrites) + @classmethod + def _apply_guardrail_tool_calls_to_output( + cls, + response_content: Sequence[object], + scanned_tool_calls: Sequence[object], + pre_guardrail_tool_calls: tuple[_ToolCallShape, ...], + returned_tool_calls: Sequence[object] | None, + guardrail_name: str | None, + ) -> None: + """Land guardrail tool-call rewrites on the response's ``tool_use`` blocks, which + ``_extract_from_content_blocks`` listed in content order, so the nth returned call + belongs to the nth block. A rewrite whose arguments are not a JSON object cannot become a + ``tool_use`` input, so it fails the request instead of leaking the originals.""" + post_guardrail_tool_calls: Final = _post_guardrail_tool_call_shapes( + returned_tool_calls, pre_guardrail_tool_calls, scanned_tool_calls, guardrail_name + ) + tool_use_blocks: Final = tuple( + block + for block in response_content + if (fields := cls._output_block_fields(block)) is not None and fields[0] == "tool_use" + ) + rewritten: Final = tuple( + (block, before, after, _rewritten_tool_use_input(after.arguments)) + for block, before, after in zip(tool_use_blocks, pre_guardrail_tool_calls, post_guardrail_tool_calls) + if before != after + ) + applicable: Final = tuple( + (block, before, after, rewritten_input) + for block, before, after, rewritten_input in rewritten + if rewritten_input is not None + ) + if len(applicable) != len(rewritten): + raise unappliable_request_rewrite(guardrail_name) + for block, before, after, rewritten_input in applicable: + _write_back_tool_use_block(block, before, after, rewritten_input) + @staticmethod def _get_response_content(response: object) -> Sequence[object]: """Extract content list from a dict or object response.""" diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index 5c5204ef2f4..3505df70b0a 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -19,8 +19,9 @@ import time import uuid from collections.abc import Iterator, Mapping, Sequence from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Union, cast +from typing import TYPE_CHECKING, Final, Union, cast +from pydantic import ConfigDict, TypeAdapter, ValidationError, with_config from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm @@ -76,6 +77,84 @@ if TYPE_CHECKING: from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth +@with_config(ConfigDict(extra="ignore")) +class _ReturnedToolCallFunction(TypedDict, total=False): + name: ReadOnly[str] + arguments: ReadOnly[str] + + +@with_config(ConfigDict(extra="allow")) +class _ReturnedToolCall(TypedDict, total=False): + function: ReadOnly[_ReturnedToolCallFunction] + + +class _WritableToolCall(TypedDict, total=False): + function: ReadOnly[dict[str, object]] # mutable-ok: the request's own tool_call dict, rewritten in place + + +_RETURNED_TOOL_CALLS: Final = TypeAdapter(tuple[_ReturnedToolCall, ...]) +_MESSAGE_TOOL_CALLS: Final = TypeAdapter(tuple[object, ...]) + + +def _message_tool_calls(message: Mapping[str, object]) -> tuple[object, ...]: + """The message's own tool_call objects, so a write-back lands on the request itself.""" + try: + return _MESSAGE_TOOL_CALLS.validate_python(message.get("tool_calls")) + except ValidationError: + return () + + +def _usable_returned_tool_calls( + returned_tool_calls: object, scanned_count: int, guardrail: "CustomGuardrail" +) -> tuple[_ReturnedToolCall, ...] | None: + """The returned ``tool_calls`` when they line up one to one with the scanned calls, else ``None``, + as a list of another length cannot be matched to the calls by position.""" + try: + rows: Final = _RETURNED_TOOL_CALLS.validate_python(returned_tool_calls) + except ValidationError: + return None + if len(rows) == scanned_count: + return rows + verbose_proxy_logger.warning( + "OpenAI Chat Completions: guardrail %s returned %d tool calls for the %d scanned, leaving them unchanged", + guardrail.guardrail_name, + len(rows), + scanned_count, + ) + return None + + +def _scanned_tool_call_rows(tool_calls_to_check: Sequence[Mapping[str, object]]) -> tuple[_ReturnedToolCall, ...]: + """The scanned calls as the guardrail left them, which carry its in-place rewrites when it returned nothing.""" + try: + return _RETURNED_TOOL_CALLS.validate_python(tool_calls_to_check) + except ValidationError: + return () + + +_FUNCTION_MAPPING: Final = TypeAdapter(dict[str, object]) + + +def _function_fields(tool_call: Mapping[str, object]) -> tuple[object, object]: + try: + function: Final = _FUNCTION_MAPPING.validate_python(tool_call.get("function")) + except ValidationError: + return (None, None) + return (function.get("name"), function.get("arguments")) + + +def _write_back_tool_call_function(tool_call: object, guardrailed_tool_call: _ReturnedToolCall) -> None: + """Keep the request's own tool_call object, with its ``id`` and ``type``, and take only the + ``function`` name and arguments the guardrail handed back.""" + guardrailed_function: Final = guardrailed_tool_call.get("function") + if guardrailed_function is None or not isinstance(tool_call, dict): + return + function: Final = cast(_WritableToolCall, tool_call).get("function") # cast-ok: a message tool_call dict + if not isinstance(function, dict): + return + function.update({key: value for key, value in guardrailed_function.items() if key in ("name", "arguments")}) + + class OpenAIChatCompletionsHandler(BaseTranslation): """ Handler for processing OpenAI chat completions messages with guardrails. @@ -165,6 +244,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): inputs["model"] = model original_structured_messages: Final = inputs.get("structured_messages") + pre_guardrail_functions: Final = tuple(_function_fields(tool_call) for tool_call in tool_calls_to_check) guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail( inputs=inputs, request_data=data, @@ -173,7 +253,9 @@ class OpenAIChatCompletionsHandler(BaseTranslation): ) guardrailed_texts: Final = guardrailed_inputs.get("texts", []) - guardrailed_tool_calls: Final = guardrailed_inputs.get("tool_calls", []) + guardrailed_tool_calls: Final = _usable_returned_tool_calls( + guardrailed_inputs.get("tool_calls"), len(tool_calls_to_check), guardrail_to_apply + ) guardrailed_tools: Final = guardrailed_inputs.get("tools") if guardrailed_tools is not None: data["tools"] = ( @@ -191,7 +273,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): guardrailed_structured_messages is not None and guardrailed_structured_messages is not original_structured_messages ): - data["messages"] = ( + rewritten_messages: Final = ( guardrailed_structured_messages if guardrail_to_apply.structured_messages_cover_full_request() else merge_guardrailed_scoped_messages( @@ -200,6 +282,15 @@ class OpenAIChatCompletionsHandler(BaseTranslation): guardrailed_scoped=guardrailed_structured_messages, ) ) + data["messages"] = rewritten_messages + await self._apply_tool_calls_to_rewritten_rows( + rewritten_messages=rewritten_messages, + scanned_count=len(messages), + tool_calls=guardrailed_tool_calls, + pre_guardrail_functions=pre_guardrail_functions, + task_mappings=tool_call_task_mappings, + guardrail=guardrail_to_apply, + ) else: # Step 3: Map guardrail responses back to original message structure if guardrailed_texts and texts_to_check: @@ -212,7 +303,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): ) # Step 4: Apply guardrailed tool calls back to messages - if guardrailed_tool_calls: + if guardrailed_tool_calls is not None: await self._apply_guardrail_responses_to_input_tool_calls( messages=messages, tool_calls=guardrailed_tool_calls, @@ -363,11 +454,45 @@ class OpenAIChatCompletionsHandler(BaseTranslation): # Replace specific text item in list content content[content_idx_optional]["text"] = guardrail_response + async def _apply_tool_calls_to_rewritten_rows( + self, + rewritten_messages: Sequence[Mapping[str, object]], + scanned_count: int, + tool_calls: Sequence[_ReturnedToolCall] | None, + pre_guardrail_functions: Sequence[tuple[object, object]], + task_mappings: Sequence[tuple[int, int]], + guardrail: "CustomGuardrail", + ) -> None: + """Returned ``tool_calls`` are placed by message index, which only holds when the returned + ``structured_messages`` kept one row per row scanned. Only calls the guardrail changed are + written, so a list echoed back unchanged cannot undo a redaction the rows carry.""" + if tool_calls is None: + return + if len(rewritten_messages) != scanned_count: + verbose_proxy_logger.warning( + "OpenAI Chat Completions: guardrail %s rewrote the conversation to %d rows from %d, " + "so its tool_calls rewrite cannot be placed and is left unapplied", + guardrail.guardrail_name, + len(rewritten_messages), + scanned_count, + ) + return + changed: Final = tuple( + (mapping, tool_call) + for mapping, tool_call, before in zip(task_mappings, tool_calls, pre_guardrail_functions) + if _function_fields(tool_call) != before + ) + await self._apply_guardrail_responses_to_input_tool_calls( + messages=rewritten_messages, + tool_calls=tuple(tool_call for _, tool_call in changed), + task_mappings=tuple(mapping for mapping, _ in changed), + ) + async def _apply_guardrail_responses_to_input_tool_calls( self, messages: Sequence[Mapping[str, object]], - tool_calls: Sequence[Mapping[str, object]], - task_mappings: list[tuple[int, int]], + tool_calls: Sequence[_ReturnedToolCall], + task_mappings: Sequence[tuple[int, int]], ) -> None: """ Apply guardrailed tool calls back to input messages. @@ -378,13 +503,10 @@ class OpenAIChatCompletionsHandler(BaseTranslation): Override this method to customize how tool call responses are applied. """ for task_idx, (msg_idx, tool_call_idx) in enumerate(task_mappings): - if task_idx < len(tool_calls): - guardrailed_tool_call = tool_calls[task_idx] - message_tool_calls = messages[msg_idx].get("tool_calls", None) - if message_tool_calls is not None and isinstance(message_tool_calls, list): - if tool_call_idx < len(message_tool_calls): - # Replace the tool call with the guardrailed version - message_tool_calls[tool_call_idx] = guardrailed_tool_call + message_tool_calls = _message_tool_calls(messages[msg_idx]) + if task_idx >= len(tool_calls) or tool_call_idx >= len(message_tool_calls): + continue + _write_back_tool_call_function(message_tool_calls[tool_call_idx], tool_calls[task_idx]) async def process_output_response( self, @@ -466,12 +588,9 @@ class OpenAIChatCompletionsHandler(BaseTranslation): ) guardrailed_texts: Final = guardrailed_inputs.get("texts", []) - returned_tool_calls: Final = guardrailed_inputs.get("tool_calls") - guardrailed_tool_calls: Final[list[dict[str, object]]] = ( - cast(list[dict[str, object]], returned_tool_calls) - if isinstance(returned_tool_calls, list) and len(returned_tool_calls) == len(tool_calls_to_check) - else tool_calls_to_check - ) + guardrailed_tool_calls: Final = _usable_returned_tool_calls( + guardrailed_inputs.get("tool_calls"), len(tool_calls_to_check), guardrail_to_apply + ) or _scanned_tool_call_rows(tool_calls_to_check) # Step 3: Map guardrail responses back to original response structure if guardrailed_texts and texts_to_check: @@ -1064,7 +1183,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): async def _apply_guardrail_responses_to_output_tool_calls( self, response: "ModelResponse", - tool_calls: list[dict[str, Any]], + tool_calls: Sequence[_ReturnedToolCall], task_mappings: list[tuple[int, int]], ) -> None: """ @@ -1087,8 +1206,8 @@ class OpenAIChatCompletionsHandler(BaseTranslation): # Update the tool call with guardrailed version existing_tool_call = choice_tool_calls[tool_call_idx] # Update object attributes (output responses always have typed objects) - if "function" in guardrailed_tool_call: - func_dict = guardrailed_tool_call["function"] + func_dict = guardrailed_tool_call.get("function") + if func_dict is not None: if "arguments" in func_dict: existing_tool_call.function.arguments = func_dict["arguments"] if "name" in func_dict: 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..c18b95f602d 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 @@ -456,6 +456,8 @@ class GenericGuardrailAPI(CustomGuardrail): return_inputs["tools"] = guardrail_response.tools elif tools: return_inputs["tools"] = tools + if guardrail_response.tool_calls is not None: + return_inputs["tool_calls"] = list(guardrail_response.tool_calls) compared_rows, is_echo = ( (structured_messages, _is_echo) if self.forward_messages_as_sent else (shown_messages, _is_shown_row) ) diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py b/litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py index 4e4983c53f0..7f5b8a478fd 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py @@ -1,8 +1,16 @@ from collections.abc import Mapping, Sequence from typing import Any, Final, Literal, cast # noqa: TID251 # JSON chat rows have no typed constructor across roles -from pydantic import ConfigDict, Field, SerializerFunctionWrapHandler, model_serializer -from typing_extensions import TypedDict +from pydantic import ( + ConfigDict, + Field, + SerializerFunctionWrapHandler, + TypeAdapter, + ValidationError, + model_serializer, + with_config, +) +from typing_extensions import ReadOnly, TypedDict from litellm.types.guardrails import GuardrailNumRetries from litellm.types.llms.base import LiteLLMBaseModel @@ -188,12 +196,43 @@ def structured_messages_from_response(value: object) -> Sequence[AllMessageValue return cast("Sequence[AllMessageValues]", value) # cast-ok: JSON rows checked for a role, the same trust texts get +@with_config(ConfigDict(extra="allow")) +class _GuardrailToolCallFunction(TypedDict): + name: ReadOnly[str] + arguments: ReadOnly[str] + + +@with_config(ConfigDict(extra="allow")) +class _GuardrailToolCallRow(TypedDict): + function: ReadOnly[_GuardrailToolCallFunction] + + +_TOOL_CALL_ROWS: Final = TypeAdapter(tuple[_GuardrailToolCallRow, ...]) + + +def tool_calls_from_response(value: object) -> tuple[ChatCompletionToolCallChunk, ...] | None: + """The rewritten tool calls, or ``None`` when the answer carries none. An empty array is no rewrite + either, as the only thing it could say about the calls it received is to drop them, which a + positional rewrite cannot do, and forwarding it would read as a count change downstream.""" + try: + rows: Final = _TOOL_CALL_ROWS.validate_python(value) + except ValidationError: + return None + if not rows: + return None + validated: Final = cast( # cast-ok: rows validated to carry function.name/arguments + tuple[ChatCompletionToolCallChunk, ...], rows + ) + return validated + + class GenericGuardrailAPIResponse: """Response model for the Generic Guardrail API""" texts: list[str] | None images: list[str] | None tools: list[GuardrailToolParam] | None + tool_calls: Sequence[ChatCompletionToolCallChunk] | None structured_messages: Sequence[AllMessageValues] | None action: str blocked_reason: str | None @@ -208,12 +247,14 @@ class GenericGuardrailAPIResponse: tools: list[GuardrailToolParam] | None = None, stream_holdback_chars: list[int] | None = None, structured_messages: Sequence[AllMessageValues] | None = None, + tool_calls: Sequence[ChatCompletionToolCallChunk] | None = None, ) -> None: self.action = action self.blocked_reason = blocked_reason self.texts = texts self.images = images self.tools = tools + self.tool_calls = tool_calls self.structured_messages = structured_messages # Number of trailing chars, indexed the same as ``texts``, that the # framework must withhold from streaming emission until the next @@ -226,6 +267,7 @@ class GenericGuardrailAPIResponse: stream_holdback_chars: Final = ( [coerce_stream_holdback_value(value) for value in raw_holdback] if isinstance(raw_holdback, list) else None ) + fields: Final = cast(Mapping[str, object], data) # cast-ok: the decoded JSON body the guardrail answered with return cls( action=data.get("action", "NONE"), blocked_reason=data.get("blocked_reason"), @@ -234,4 +276,5 @@ class GenericGuardrailAPIResponse: tools=data.get("tools"), stream_holdback_chars=stream_holdback_chars, structured_messages=structured_messages_from_response(data.get("structured_messages")), + tool_calls=tool_calls_from_response(fields.get("tool_calls")), ) diff --git a/tests/unit/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py b/tests/unit/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py index c2ee58b67ed..fc27f6f302f 100644 --- a/tests/unit/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py +++ b/tests/unit/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py @@ -2836,6 +2836,236 @@ class TestAnthropicMessagesTopLevelSystemAndToolUseInputs: assert data["system"] == original["system"], "a rejected rewrite must leave the request untouched" assert data["messages"] == original["messages"], "a rejected rewrite must leave the request untouched" + @staticmethod + def _tool_use_response(*tool_inputs: dict[str, object]) -> dict[str, Any]: + return { + "id": "msg_01", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-5", + "content": [ + {"type": "text", "text": "Running the checks now."}, + *( + {"type": "tool_use", "id": f"toolu_{idx:02d}", "name": "Bash", "input": tool_input} + for idx, tool_input in enumerate(tool_inputs, start=1) + ), + ], + "stop_reason": "tool_use", + "stop_sequence": None, + "usage": {"input_tokens": 10, "output_tokens": 10}, + } + + @pytest.mark.asyncio + @pytest.mark.parametrize("return_copies", [False, True]) + async def test_masked_output_tool_call_arguments_write_back_into_the_tool_use_input(self, return_copies: bool): + handler = AnthropicMessagesHandler() + guardrail = ToolCallArgumentsMaskingGuardrail(return_copies=return_copies) + response = self._tool_use_response({"cmd": "echo safe"}, {"cmd": "curl -H 'x-token: POISON' example.com"}) + + await handler.process_output_response(response=response, guardrail_to_apply=guardrail) + + assert [tool_call["function"]["name"] for tool_call in guardrail.seen_tool_calls] == ["Bash", "Bash"] + assert response["content"] == [ + {"type": "text", "text": "Running the checks now."}, + {"type": "tool_use", "id": "toolu_01", "name": "Bash", "input": {"cmd": "echo safe"}}, + { + "type": "tool_use", + "id": "toolu_02", + "name": "Bash", + "input": {"cmd": "curl -H 'x-token: [BLOCKED]' example.com"}, + }, + ] + + @pytest.mark.asyncio + async def test_masked_output_tool_call_arguments_write_back_into_object_shaped_tool_use_blocks(self): + from pydantic import BaseModel + + class _TextBlock(BaseModel): + type: str = "text" + text: str + + class _ToolUseBlock(BaseModel): + type: str = "tool_use" + id: str + name: str + input: dict[str, object] + + class _Response(BaseModel): + content: list[_TextBlock | _ToolUseBlock] + + handler = AnthropicMessagesHandler() + guardrail = ToolCallArgumentsMaskingGuardrail() + response = _Response( + content=[ + _TextBlock(text="Running the checks now."), + _ToolUseBlock(id="toolu_01", name="Bash", input={"cmd": "curl -H 'x-token: POISON' example.com"}), + ] + ) + + await handler.process_output_response(response=response, guardrail_to_apply=guardrail) + + assert [tool_call["function"]["name"] for tool_call in guardrail.seen_tool_calls] == ["Bash"] + assert response.content[1].input == {"cmd": "curl -H 'x-token: [BLOCKED]' example.com"}, ( + "a rewrite computed for an object-shaped tool_use block must land on that block" + ) + + @pytest.mark.asyncio + async def test_a_renamed_output_tool_call_writes_back_the_tool_use_name(self): + class RenamingGuardrail(ToolCallArgumentsMaskingGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict[str, object], + input_type: Literal["request", "response"], + logging_obj: Optional[LiteLLMLoggingObj] = None, + ) -> GenericGuardrailAPIInputs: + outputs = await super().apply_guardrail(inputs, request_data, input_type, logging_obj) + for tool_call in outputs.get("tool_calls") or []: + tool_call["function"]["name"] = "SafeBash" + return outputs + + handler = AnthropicMessagesHandler() + response = self._tool_use_response({"cmd": "echo POISON"}) + + await handler.process_output_response( + response=response, guardrail_to_apply=RenamingGuardrail(return_copies=True) + ) + + assert response["content"][1] == { + "type": "tool_use", + "id": "toolu_01", + "name": "SafeBash", + "input": {"cmd": "echo [BLOCKED]"}, + } + + @pytest.mark.asyncio + async def test_non_json_rewritten_output_arguments_are_rejected_by_name(self): + from litellm.llms.base_llm.guardrail_translation.utils import UnappliableRequestRewrite + + handler = AnthropicMessagesHandler() + guardrail = ToolCallArgumentsMaskingGuardrail(replacement_arguments="[REDACTED]") + response = self._tool_use_response({"cmd": "echo POISON"}) + original = json.loads(json.dumps(response)) + + with pytest.raises(UnappliableRequestRewrite) as excinfo: + await handler.process_output_response(response=response, guardrail_to_apply=guardrail) + + assert excinfo.value.guardrail_name == "scan-only-capture" + assert response == original, "a rejected rewrite must leave the response untouched" + + @staticmethod + def _tool_call_list_of_another_length( + returned: Literal["drops_first", "empty", "doubled"], + ) -> type[ToolCallArgumentsMaskingGuardrail]: + class AnotherLengthGuardrail(ToolCallArgumentsMaskingGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict[str, object], + input_type: Literal["request", "response"], + logging_obj: Optional[LiteLLMLoggingObj] = None, + ) -> GenericGuardrailAPIInputs: + outputs = await super().apply_guardrail(inputs, request_data, input_type, logging_obj) + masked = list(outputs.get("tool_calls") or []) + outputs["tool_calls"] = {"drops_first": masked[1:], "empty": [], "doubled": masked * 2}[returned] + return outputs + + return AnotherLengthGuardrail + + @pytest.mark.asyncio + @pytest.mark.parametrize("returned", ["drops_first", "empty", "doubled"]) + async def test_an_output_tool_call_list_of_another_length_leaves_the_tool_use_blocks_as_scanned( + self, returned: Literal["drops_first", "empty", "doubled"] + ): + handler = AnthropicMessagesHandler() + guardrail = self._tool_call_list_of_another_length(returned)(return_copies=True) + response = self._tool_use_response({"cmd": "echo POISON"}, {"cmd": "echo POISON twice"}) + original = json.loads(json.dumps(response)) + + processed = await handler.process_output_response(response=response, guardrail_to_apply=guardrail) + + assert processed is response + assert response == original, "a tool_calls list that does not line up with the blocks is left unapplied" + + @pytest.mark.asyncio + async def test_an_input_tool_call_list_of_another_length_leaves_the_tool_use_input_as_scanned(self): + handler = AnthropicMessagesHandler() + guardrail = self._tool_call_list_of_another_length("empty")(return_copies=True) + data = self._tool_use_conversation(system="You are a careful agent harness.") + original = json.loads(json.dumps(data)) + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert data["messages"] == original["messages"], "an empty tool_calls list is not a rewrite of the calls" + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "leading_blocks", + [ + [], + [{"type": "text", "text": "Let me check."}, {"type": "text", "text": "Running it now."}], + ], + ids=["tool_use_only", "two_text_blocks_before_tool_use"], + ) + async def test_tool_call_rewrites_land_when_the_guardrail_also_rewrites_the_message_rows(self, leading_blocks): + class RowsAndToolCallsGuardrail(MockStructuredMaskingGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict[str, object], + input_type: Literal["request", "response"], + logging_obj: Optional[LiteLLMLoggingObj] = None, + ) -> GenericGuardrailAPIInputs: + outputs = await super().apply_guardrail(inputs, request_data, input_type, logging_obj) + outputs["tool_calls"] = [ + { + **tool_call, + "function": { + **tool_call["function"], + "arguments": tool_call["function"]["arguments"].replace("POISON", "[BLOCKED]"), + }, + } + for tool_call in inputs.get("tool_calls") or [] + ] + return outputs + + handler = AnthropicMessagesHandler() + data = self._tool_use_conversation(system="You are a careful agent harness.") + data["messages"][0]["content"] = "run the check for bob@example.com" + data["messages"][1]["content"] = [*leading_blocks, *data["messages"][1]["content"]] + + await handler.process_input_messages(data=data, guardrail_to_apply=RowsAndToolCallsGuardrail()) + + first_turn = json.dumps(data["messages"][0]["content"]) + assert "" in first_turn and "bob@example.com" not in first_turn, "the row rewrite must land" + tool_use_blocks = [b for b in data["messages"][1]["content"] if b.get("type") == "tool_use"] + assert [b["input"] for b in tool_use_blocks] == [ + {"cmd": "AWS_ACCESS_KEY_ID=[BLOCKED] aws sts get-caller-identity"} + ], "the tool_calls rewrite must land on the rewritten rows too" + assert data["messages"][2]["content"][0]["tool_use_id"] == "toolu_01" + + @pytest.mark.asyncio + async def test_output_arguments_rewritten_to_an_empty_string_become_an_empty_tool_use_input(self): + handler = AnthropicMessagesHandler() + guardrail = ToolCallArgumentsMaskingGuardrail(return_copies=True, replacement_arguments="") + response = self._tool_use_response({"cmd": "echo POISON"}) + + await handler.process_output_response(response=response, guardrail_to_apply=guardrail) + + assert response["content"][1] == {"type": "tool_use", "id": "toolu_01", "name": "Bash", "input": {}} + assert json.loads(json.dumps(response)) == response, "the rewritten input must serialize like any other block" + + @pytest.mark.asyncio + async def test_input_arguments_rewritten_to_an_empty_string_become_an_empty_tool_use_input(self): + handler = AnthropicMessagesHandler() + guardrail = ToolCallArgumentsMaskingGuardrail(return_copies=True, replacement_arguments="") + data = self._tool_use_conversation(system="You are a careful agent harness.") + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert data["messages"][1]["content"][0] == {"type": "tool_use", "id": "toolu_01", "name": "Bash", "input": {}} + assert json.loads(json.dumps(data["messages"])) == data["messages"], "the request must still serialize" + @pytest.mark.asyncio async def test_scan_only_tool_results_keeps_system_and_tool_use_out(self): handler = AnthropicMessagesHandler() diff --git a/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py b/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py index 5d71cbe9afc..6ca310e1d35 100644 --- a/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py +++ b/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py @@ -907,6 +907,151 @@ class TestOpenAIChatCompletionsHandlerToolCallsOutput: assert tool_calls[1].function.arguments == second_args +class MockFunctionOnlyToolCallGuardrail(CustomGuardrail): + """Mock guardrail that returns each tool call as only its rewritten function, the way a + scrubber that does not echo id and type does.""" + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + return GenericGuardrailAPIInputs( + texts=inputs.get("texts", []), + tool_calls=[ + {"function": {"name": tool_call["function"]["name"], "arguments": json.dumps({"email": "[EMAIL]"})}} + for tool_call in inputs.get("tool_calls", []) + ], + ) + + +class MockRowsAndToolCallsGuardrail(CustomGuardrail): + """Mock guardrail that rewrites the message rows and the tool calls in one answer.""" + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + return GenericGuardrailAPIInputs( + texts=[text.replace("john@example.com", "[EMAIL]") for text in inputs.get("texts", [])], + structured_messages=[ + {**message, "content": message["content"].replace("john@example.com", "[EMAIL]")} + if isinstance(message.get("content"), str) + else dict(message) + for message in inputs.get("structured_messages") or [] + ], + tool_calls=[ + {**tool_call, "function": {**tool_call["function"], "arguments": json.dumps({"email": "[EMAIL]"})}} + for tool_call in inputs.get("tool_calls", []) + ], + ) + + +class MockRowRedactingEchoingToolCallsGuardrail(CustomGuardrail): + """Mock guardrail that redacts a tool call inside the message rows and echoes the scanned + tool_calls list back unchanged, the way a backend that only answers rows does.""" + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + return GenericGuardrailAPIInputs( + texts=inputs.get("texts", []), + structured_messages=[ + json.loads(json.dumps(message).replace("john@example.com", "[EMAIL]")) + for message in inputs.get("structured_messages") or [] + ], + tool_calls=list(inputs.get("tool_calls", [])), + ) + + +class TestOpenAIChatCompletionsHandlerToolCallsInputWriteBack: + @pytest.mark.asyncio + async def test_echoed_tool_calls_do_not_undo_a_redaction_made_in_the_message_rows(self): + handler = OpenAIChatCompletionsHandler() + data = self._assistant_tool_call_turn(json.dumps({"email": "john@example.com"})) + + await handler.process_input_messages( + data, MockRowRedactingEchoingToolCallsGuardrail(guardrail_name="test") + ) + + assert data["messages"][1]["tool_calls"][0]["function"]["arguments"] == json.dumps({"email": "[EMAIL]"}), ( + "an echoed, unchanged tool_calls list must not overwrite the redaction the rows carry" + ) + + @pytest.mark.asyncio + async def test_tool_call_rewrites_land_when_the_guardrail_also_rewrites_the_message_rows(self): + handler = OpenAIChatCompletionsHandler() + data = self._assistant_tool_call_turn(json.dumps({"email": "john@example.com"})) + data["messages"][0]["content"] = "email john@example.com" + + await handler.process_input_messages(data, MockRowsAndToolCallsGuardrail(guardrail_name="test")) + + assert data["messages"][0]["content"] == "email [EMAIL]", "the row rewrite must land" + assert data["messages"][1]["tool_calls"] == [ + { + "id": "call_1", + "type": "function", + "function": {"name": "send_email", "arguments": json.dumps({"email": "[EMAIL]"})}, + } + ], "the tool_calls rewrite must land on the rewritten rows too" + + @staticmethod + def _assistant_tool_call_turn(*arguments: str) -> dict[str, Any]: + return { + "messages": [ + {"role": "user", "content": "email them"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": f"call_{idx}", + "type": "function", + "function": {"name": "send_email", "arguments": call_arguments}, + } + for idx, call_arguments in enumerate(arguments, start=1) + ], + }, + ] + } + + @pytest.mark.asyncio + async def test_returned_function_only_rows_keep_the_request_tool_call_id_and_type(self): + handler = OpenAIChatCompletionsHandler() + data = self._assistant_tool_call_turn(json.dumps({"email": "john@example.com"})) + + await handler.process_input_messages(data, MockFunctionOnlyToolCallGuardrail(guardrail_name="test")) + + assert data["messages"][1]["tool_calls"] == [ + { + "id": "call_1", + "type": "function", + "function": {"name": "send_email", "arguments": json.dumps({"email": "[EMAIL]"})}, + } + ] + + @pytest.mark.asyncio + async def test_a_returned_list_of_another_length_leaves_the_request_tool_calls_unchanged(self): + handler = OpenAIChatCompletionsHandler() + first_args = json.dumps({"email": "a@example.com"}) + second_args = json.dumps({"email": "b@example.com"}) + data = self._assistant_tool_call_turn(first_args, second_args) + original = json.loads(json.dumps(data)) + + await handler.process_input_messages(data, MockMisalignedToolCallGuardrail(guardrail_name="test")) + + assert data["messages"] == original["messages"], "one returned call for two scanned cannot be placed" + + class MockPassThroughGuardrail(CustomGuardrail): """Mock guardrail that passes through without blocking - for testing streaming fallback behavior""" diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py index b5f794cb136..392367036e2 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py @@ -2692,3 +2692,63 @@ async def test_during_mcp_call_blocked_raises() -> None: ) assert exc_info.value.message == "credential detected in tool input" + + +class TestToolCallsInResponse: + @staticmethod + def _tool_calls(arguments: str) -> list[dict[str, object]]: + return [{"id": "call_1", "type": "function", "function": {"name": "run_shell", "arguments": arguments}}] + + @staticmethod + def _guardrail_answering(body: dict[str, object], sent_bodies: list[dict[str, object]]) -> GenericGuardrailAPI: + async def answer(request: httpx.Request) -> httpx.Response: + sent_bodies.append(json.loads(request.content)) + return httpx.Response(200, json=body, request=request) + + return GenericGuardrailAPI( + api_base="https://api.test.guardrail.com", + guardrail_name="tool-call-rewriter", + event_hook="post_call", + default_on=True, + async_handler=AsyncHTTPHandler(transport=httpx.MockTransport(answer)), + ) + + @pytest.mark.asyncio + async def test_rewritten_tool_calls_from_the_guardrail_reach_the_caller(self, mock_request_data_input): + rewritten: Final = self._tool_calls('{"cmd": "echo [REDACTED]"}') + sent_bodies: Final[list[dict[str, object]]] = [] # mutable-ok: records the body the guardrail endpoint received + guardrail: Final = self._guardrail_answering( + {"action": "GUARDRAIL_INTERVENED", "texts": [], "tool_calls": rewritten}, sent_bodies + ) + + guardrailed_inputs: Final = await guardrail.apply_guardrail( + inputs={"texts": [], "tool_calls": self._tool_calls('{"cmd": "echo SECRET-123"}')}, + request_data=mock_request_data_input, + input_type="response", + ) + + assert sent_bodies[0]["tool_calls"] == self._tool_calls('{"cmd": "echo SECRET-123"}') + assert guardrailed_inputs["tool_calls"] == rewritten + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "body", + [ + {"action": "NONE", "texts": []}, + {"action": "NONE", "texts": [], "tool_calls": "not a list"}, + {"action": "NONE", "texts": [], "tool_calls": ["not a tool call"]}, + {"action": "NONE", "texts": [], "tool_calls": [{"id": "call_1", "type": "function", "function": None}]}, + {"action": "NONE", "texts": [], "tool_calls": [{"function": {"name": "run_shell"}}]}, + {"action": "GUARDRAIL_INTERVENED", "texts": ["my ssn is [MASKED]"], "tool_calls": []}, + ], + ) + async def test_an_answer_without_usable_tool_calls_leaves_the_key_out(self, mock_request_data_input, body): + guardrail: Final = self._guardrail_answering(body, []) + + guardrailed_inputs: Final = await guardrail.apply_guardrail( + inputs={"texts": [], "tool_calls": self._tool_calls('{"cmd": "echo SECRET-123"}')}, + request_data=mock_request_data_input, + input_type="response", + ) + + assert "tool_calls" not in guardrailed_inputs