mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-14 09:03:16 +00:00
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 <ivana.dunisijevic@mistral.ai>
This commit is contained in:
jesus
parent
3b35ae340d
commit
b8e7f13aeb
7 changed files with 802 additions and 34 deletions
|
|
@ -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."""
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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")),
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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 "<EMAIL>" 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()
|
||||
|
|
|
|||
|
|
@ -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"""
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue