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 guardrail's response tool_calls were dropped at parse time and the /v1/messages output handler never applied them, so a post_call guardrail that masks tool-call arguments had no effect on any endpoint. The hook now parses tool_calls, matches them one to one with the calls it sent, keeps each call's own id and type, and forwards them. An empty array is no rewrite. A list of another length or a malformed entry is a guardrail error governed by fail_on_error. The Anthropic output handler writes the returned calls onto the tool_use blocks by position.
This commit is contained in:
Jesus Velarde
parent
bbc7de1ccb
commit
8ea78482cc
5 changed files with 512 additions and 33 deletions
|
|
@ -156,6 +156,12 @@ class _TextAttrBlock(Protocol):
|
|||
text: str
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class _ToolUseAttrBlock(Protocol):
|
||||
name: str
|
||||
input: Mapping[str, object]
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class _WritableMessage(Protocol):
|
||||
@overload
|
||||
|
|
@ -443,13 +449,23 @@ def _rewritten_tool_use_input(arguments: str) -> Mapping[str, object] | None:
|
|||
def _write_back_tool_use(
|
||||
message: _WritableMessage, target: ToolUseInputTarget, shape: _ToolCallShape, rewritten_input: Mapping[str, object]
|
||||
) -> None:
|
||||
content: Final = message.get("content", None)
|
||||
block: Final = content[target.content_idx] if isinstance(content, list) else None
|
||||
if not isinstance(block, dict):
|
||||
blocks: Final = list_field(message, "content")
|
||||
_write_back_tool_use_block(
|
||||
blocks[target.content_idx] if target.content_idx < len(blocks) else None, shape, rewritten_input
|
||||
)
|
||||
|
||||
|
||||
def _write_back_tool_use_block(block: object, shape: _ToolCallShape, rewritten_input: Mapping[str, object]) -> None:
|
||||
if isinstance(block, dict):
|
||||
block["input"] = rewritten_input # rebind-ok: the tool_use block is the client-visible object being rewritten
|
||||
if shape.name is not None and shape.name != block.get("name"):
|
||||
block["name"] = shape.name # rebind-ok: same in-place rewrite as the input above
|
||||
return
|
||||
block["input"] = rewritten_input
|
||||
if shape.name is not None and shape.name != block.get("name"):
|
||||
block["name"] = shape.name
|
||||
if not isinstance(block, _ToolUseAttrBlock):
|
||||
return
|
||||
block.input = rewritten_input # rebind-ok: same in-place rewrite for an object-shaped tool_use block
|
||||
if shape.name is not None and shape.name != block.name:
|
||||
block.name = shape.name # rebind-ok: same in-place rewrite as the input above
|
||||
|
||||
|
||||
def _streamed_query_strings(
|
||||
|
|
@ -1520,6 +1536,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,
|
||||
|
|
@ -1541,6 +1558,13 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
self._write_back_web_search_strings(
|
||||
response_content, _web_search_rewrites_by_block(web_search_strings, web_search_rewrites)
|
||||
)
|
||||
self._write_back_output_tool_calls(
|
||||
response_content,
|
||||
scanned_tool_calls=tool_calls_to_check,
|
||||
pre_guardrail_tool_calls=pre_guardrail_tool_calls,
|
||||
returned_tool_calls=guardrailed_inputs.get("tool_calls"),
|
||||
guardrail_name=guardrail_to_apply.guardrail_name,
|
||||
)
|
||||
|
||||
verbose_proxy_logger.debug("Anthropic Messages: Processed output response: %s", response)
|
||||
|
||||
|
|
@ -1717,6 +1741,43 @@ 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 _write_back_output_tool_calls(
|
||||
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:
|
||||
"""The nth returned call belongs to the nth ``tool_use`` block, the order the calls were scanned
|
||||
in. A returned list of another length cannot be placed, so the scanned calls stand, as on the
|
||||
request side. Rewritten arguments that are not a JSON object cannot become a ``tool_use`` input,
|
||||
so the response fails instead of carrying the arguments the guardrail meant to replace."""
|
||||
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 scanned_tool_calls
|
||||
)
|
||||
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"
|
||||
)
|
||||
if len(tool_use_blocks) != len(pre_guardrail_tool_calls):
|
||||
return
|
||||
rewrites: Final = tuple(
|
||||
(block, 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
|
||||
)
|
||||
if any(rewritten_input is None for _, _, rewritten_input in rewrites):
|
||||
raise unappliable_request_rewrite(guardrail_name)
|
||||
for block, after, rewritten_input in rewrites:
|
||||
if rewritten_input is not None:
|
||||
_write_back_tool_use_block(block, after, rewritten_input)
|
||||
|
||||
@staticmethod
|
||||
def _get_response_content(response: object) -> Sequence[object]:
|
||||
"""Extract content list from a dict or object response."""
|
||||
|
|
|
|||
|
|
@ -34,6 +34,7 @@ from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import
|
|||
GenericGuardrailAPIRequest,
|
||||
GenericGuardrailAPIResponse,
|
||||
GuardrailToolParam,
|
||||
tool_calls_to_write_back,
|
||||
)
|
||||
from litellm.types.utils import GenericGuardrailAPIInputs
|
||||
|
||||
|
|
@ -45,6 +46,9 @@ GUARDRAIL_NAME: Final = "generic_guardrail_api"
|
|||
_CHAT_ROW: Final[TypeAdapter[AllMessageValues]] = TypeAdapter(AllMessageValues)
|
||||
_SENT_ROW: Final = TypeAdapter(Mapping[str, object])
|
||||
_SENT_ROWS: Final = TypeAdapter(Sequence[Mapping[str, object]])
|
||||
_SENT_TOOL_CALLS: Final[TypeAdapter[tuple[Mapping[str, object], ...] | None]] = TypeAdapter(
|
||||
tuple[Mapping[str, object], ...] | None
|
||||
)
|
||||
_JSON: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue)
|
||||
|
||||
# Headers whose values are forwarded as-is (case-insensitive). Glob patterns supported (e.g. x-stainless-*, x-litellm*).
|
||||
|
|
@ -442,12 +446,16 @@ class GenericGuardrailAPI(CustomGuardrail):
|
|||
tools: list[ChatCompletionToolParam] | None,
|
||||
structured_messages: Sequence[AllMessageValues] | None,
|
||||
shown_messages: Sequence[AllMessageValues] | None,
|
||||
sent_tool_calls: Sequence[Mapping[str, object]] | None,
|
||||
guardrail_response: GenericGuardrailAPIResponse,
|
||||
) -> GenericGuardrailAPIInputs:
|
||||
# Action is NONE or no modifications needed
|
||||
return_inputs: Final = GenericGuardrailAPIInputs(texts=texts)
|
||||
if guardrail_response.texts:
|
||||
return_inputs["texts"] = guardrail_response.texts
|
||||
rewritten_tool_calls: Final = tool_calls_to_write_back(sent_tool_calls, guardrail_response.tool_calls)
|
||||
if rewritten_tool_calls is not None:
|
||||
return_inputs["tool_calls"] = list(rewritten_tool_calls)
|
||||
if guardrail_response.images:
|
||||
return_inputs["images"] = guardrail_response.images
|
||||
elif images:
|
||||
|
|
@ -601,11 +609,8 @@ class GenericGuardrailAPI(CustomGuardrail):
|
|||
|
||||
# Make the API request
|
||||
# Use mode="json" to ensure all iterables are converted to lists
|
||||
response: Final = await self._post_check(
|
||||
self._request_body(guardrail_request, structured_messages),
|
||||
headers,
|
||||
self.num_retries,
|
||||
)
|
||||
guardrail_request_body: Final = self._request_body(guardrail_request, structured_messages)
|
||||
response: Final = await self._post_check(guardrail_request_body, headers, self.num_retries)
|
||||
|
||||
response.raise_for_status()
|
||||
response_json: Final = response.json()
|
||||
|
|
@ -632,6 +637,7 @@ class GenericGuardrailAPI(CustomGuardrail):
|
|||
tools=tools,
|
||||
structured_messages=structured_messages,
|
||||
shown_messages=guardrail_request.structured_messages,
|
||||
sent_tool_calls=_SENT_TOOL_CALLS.validate_python(guardrail_request_body.get("tool_calls")),
|
||||
guardrail_response=guardrail_response,
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -1,8 +1,8 @@
|
|||
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, 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 +188,65 @@ 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 GuardrailToolCall(TypedDict):
|
||||
"""A rewritten tool call, one per call the guardrail received. The call's own ``id`` and
|
||||
``type`` are kept by the proxy, so ``function`` is all a guardrail has to send back."""
|
||||
|
||||
function: ReadOnly[GuardrailToolCallFunction]
|
||||
|
||||
|
||||
_TOOL_CALLS: Final = TypeAdapter(tuple[GuardrailToolCall, ...])
|
||||
_SENT_FUNCTION: Final = TypeAdapter(Mapping[str, object])
|
||||
_RESPONSE_FIELDS: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object])
|
||||
|
||||
|
||||
def tool_calls_from_response(value: object) -> tuple[GuardrailToolCall, ...] | None:
|
||||
"""A malformed ``tool_calls`` raises, so the guardrail's ``fail_on_error`` setting decides the
|
||||
request's fate instead of the unrewritten calls going through."""
|
||||
return None if value is None else _TOOL_CALLS.validate_python(value)
|
||||
|
||||
|
||||
def _with_returned_function(sent_call: Mapping[str, object], returned_call: GuardrailToolCall) -> Mapping[str, object]:
|
||||
sent_function: Final = _SENT_FUNCTION.validate_python(sent_call.get("function") or {})
|
||||
returned_function: Final = returned_call["function"]
|
||||
return {
|
||||
**sent_call,
|
||||
"function": {**sent_function, "name": returned_function["name"], "arguments": returned_function["arguments"]},
|
||||
}
|
||||
|
||||
|
||||
def tool_calls_to_write_back(
|
||||
sent_calls: Sequence[Mapping[str, object]] | None, returned_calls: Sequence[GuardrailToolCall] | None
|
||||
) -> tuple[ChatCompletionToolCallChunk, ...] | None:
|
||||
"""The calls as the guardrail received them, carrying the ``function`` it handed back, so every
|
||||
endpoint keeps each call's own ``id`` and ``type``. The guardrail answers one call per call received,
|
||||
in the same order, and an empty list is no rewrite. Any other length cannot be matched to the calls
|
||||
by position, which is a guardrail error rather than a reason to send the calls unrewritten."""
|
||||
if not returned_calls:
|
||||
return None
|
||||
sent_count: Final = 0 if sent_calls is None else len(sent_calls)
|
||||
if sent_calls is None or len(returned_calls) != sent_count:
|
||||
raise ValueError(f"the guardrail returned {len(returned_calls)} tool_calls for the {sent_count} it received")
|
||||
return cast( # cast-ok: the rows the proxy sent as tool calls, each with a pydantic-validated function
|
||||
tuple[ChatCompletionToolCallChunk, ...],
|
||||
tuple(_with_returned_function(sent, returned) for sent, returned in zip(sent_calls, returned_calls)),
|
||||
)
|
||||
|
||||
|
||||
class GenericGuardrailAPIResponse:
|
||||
"""Response model for the Generic Guardrail API"""
|
||||
|
||||
texts: list[str] | None
|
||||
images: list[str] | None
|
||||
tools: list[GuardrailToolParam] | None
|
||||
tool_calls: Sequence[GuardrailToolCall] | None
|
||||
structured_messages: Sequence[AllMessageValues] | None
|
||||
action: str
|
||||
blocked_reason: str | None
|
||||
|
|
@ -208,12 +261,14 @@ class GenericGuardrailAPIResponse:
|
|||
tools: list[GuardrailToolParam] | None = None,
|
||||
stream_holdback_chars: list[int] | None = None,
|
||||
structured_messages: Sequence[AllMessageValues] | None = None,
|
||||
tool_calls: Sequence[GuardrailToolCall] | 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 +281,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 = _RESPONSE_FIELDS.validate_python(data)
|
||||
return cls(
|
||||
action=data.get("action", "NONE"),
|
||||
blocked_reason=data.get("blocked_reason"),
|
||||
|
|
@ -234,4 +290,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")),
|
||||
)
|
||||
|
|
|
|||
|
|
@ -8,7 +8,8 @@ with guardrail transformations, specifically testing edge cases with empty choic
|
|||
import json
|
||||
import logging
|
||||
import re
|
||||
from typing import Any, Literal, Optional, Final
|
||||
from collections.abc import Callable, Sequence
|
||||
from typing import Any, Literal, Optional, Final, cast
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
|
@ -21,6 +22,8 @@ from litellm.llms.anthropic.chat.guardrail_translation.handler import (
|
|||
AnthropicMessagesHandler,
|
||||
)
|
||||
from litellm.llms.anthropic.messages.adapters.transformation import truncate_tool_name
|
||||
from litellm.llms.base_llm.guardrail_translation.utils import UnappliableRequestRewrite
|
||||
from litellm.types.llms.anthropic import AnthropicResponseContentBlockText, AnthropicResponseContentBlockToolUse
|
||||
from litellm.types.utils import ChatCompletionMessageToolCall, GenericGuardrailAPIInputs
|
||||
|
||||
|
||||
|
|
@ -274,12 +277,42 @@ class TestAnthropicMessagesHandlerStreamingOutputProcessing:
|
|||
@staticmethod
|
||||
def _ended_sse_chunks() -> list:
|
||||
events = [
|
||||
("message_start", {"type": "message_start", "message": {"id": "msg_1", "type": "message", "role": "assistant", "model": "claude-sonnet-4-5", "content": [], "stop_reason": None, "usage": {"input_tokens": 1, "output_tokens": 0}}}),
|
||||
("content_block_start", {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}),
|
||||
("content_block_delta", {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "hello "}}),
|
||||
("content_block_delta", {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "world"}}),
|
||||
(
|
||||
"message_start",
|
||||
{
|
||||
"type": "message_start",
|
||||
"message": {
|
||||
"id": "msg_1",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": "claude-sonnet-4-5",
|
||||
"content": [],
|
||||
"stop_reason": None,
|
||||
"usage": {"input_tokens": 1, "output_tokens": 0},
|
||||
},
|
||||
},
|
||||
),
|
||||
(
|
||||
"content_block_start",
|
||||
{"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}},
|
||||
),
|
||||
(
|
||||
"content_block_delta",
|
||||
{"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "hello "}},
|
||||
),
|
||||
(
|
||||
"content_block_delta",
|
||||
{"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "world"}},
|
||||
),
|
||||
("content_block_stop", {"type": "content_block_stop", "index": 0}),
|
||||
("message_delta", {"type": "message_delta", "delta": {"stop_reason": "end_turn", "stop_sequence": None}, "usage": {"output_tokens": 2}}),
|
||||
(
|
||||
"message_delta",
|
||||
{
|
||||
"type": "message_delta",
|
||||
"delta": {"stop_reason": "end_turn", "stop_sequence": None},
|
||||
"usage": {"output_tokens": 2},
|
||||
},
|
||||
),
|
||||
("message_stop", {"type": "message_stop"}),
|
||||
]
|
||||
return [f"event: {name}\ndata: {json.dumps(payload)}\n\n".encode() for name, payload in events]
|
||||
|
|
@ -376,13 +409,58 @@ class TestAnthropicMessagesHandlerStreamingOutputProcessing:
|
|||
@staticmethod
|
||||
def _ended_tool_use_sse_chunks() -> list:
|
||||
events = [
|
||||
("message_start", {"type": "message_start", "message": {"id": "msg_1", "type": "message", "role": "assistant", "model": "claude-sonnet-4-5", "content": [], "stop_reason": None, "usage": {"input_tokens": 1, "output_tokens": 0}}}),
|
||||
("content_block_start", {"type": "content_block_start", "index": 0, "content_block": {"type": "tool_use", "id": "toolu_1", "name": "lookup_fruit", "input": {}}}),
|
||||
("content_block_delta", {"type": "content_block_delta", "index": 0, "delta": {"type": "input_json_delta", "partial_json": ""}}),
|
||||
("content_block_delta", {"type": "content_block_delta", "index": 0, "delta": {"type": "input_json_delta", "partial_json": '{"fruit": "persim'}}),
|
||||
("content_block_delta", {"type": "content_block_delta", "index": 0, "delta": {"type": "input_json_delta", "partial_json": 'mon"}'}}),
|
||||
(
|
||||
"message_start",
|
||||
{
|
||||
"type": "message_start",
|
||||
"message": {
|
||||
"id": "msg_1",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": "claude-sonnet-4-5",
|
||||
"content": [],
|
||||
"stop_reason": None,
|
||||
"usage": {"input_tokens": 1, "output_tokens": 0},
|
||||
},
|
||||
},
|
||||
),
|
||||
(
|
||||
"content_block_start",
|
||||
{
|
||||
"type": "content_block_start",
|
||||
"index": 0,
|
||||
"content_block": {"type": "tool_use", "id": "toolu_1", "name": "lookup_fruit", "input": {}},
|
||||
},
|
||||
),
|
||||
(
|
||||
"content_block_delta",
|
||||
{"type": "content_block_delta", "index": 0, "delta": {"type": "input_json_delta", "partial_json": ""}},
|
||||
),
|
||||
(
|
||||
"content_block_delta",
|
||||
{
|
||||
"type": "content_block_delta",
|
||||
"index": 0,
|
||||
"delta": {"type": "input_json_delta", "partial_json": '{"fruit": "persim'},
|
||||
},
|
||||
),
|
||||
(
|
||||
"content_block_delta",
|
||||
{
|
||||
"type": "content_block_delta",
|
||||
"index": 0,
|
||||
"delta": {"type": "input_json_delta", "partial_json": 'mon"}'},
|
||||
},
|
||||
),
|
||||
("content_block_stop", {"type": "content_block_stop", "index": 0}),
|
||||
("message_delta", {"type": "message_delta", "delta": {"stop_reason": "tool_use", "stop_sequence": None}, "usage": {"output_tokens": 2}}),
|
||||
(
|
||||
"message_delta",
|
||||
{
|
||||
"type": "message_delta",
|
||||
"delta": {"stop_reason": "tool_use", "stop_sequence": None},
|
||||
"usage": {"output_tokens": 2},
|
||||
},
|
||||
),
|
||||
("message_stop", {"type": "message_stop"}),
|
||||
]
|
||||
return [f"event: {name}\ndata: {json.dumps(payload)}\n\n".encode() for name, payload in events]
|
||||
|
|
@ -403,7 +481,8 @@ class TestAnthropicMessagesHandlerStreamingOutputProcessing:
|
|||
json.loads(line[len("data:") :].strip())["delta"]["partial_json"]
|
||||
for chunk in chunks
|
||||
for line in chunk.decode().split("\n")
|
||||
if line.startswith("data:") and json.loads(line[len("data:") :].strip()).get("type") == "content_block_delta"
|
||||
if line.startswith("data:")
|
||||
and json.loads(line[len("data:") :].strip()).get("type") == "content_block_delta"
|
||||
]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -496,7 +575,10 @@ class TestAnthropicMessagesHandlerStreamingOutputProcessing:
|
|||
logging_obj: Optional[LiteLLMLoggingObj] = None,
|
||||
) -> GenericGuardrailAPIInputs:
|
||||
for tool_call in inputs.get("tool_calls", []):
|
||||
if isinstance(tool_call, ChatCompletionMessageToolCall) and "persimmon" in tool_call.function.arguments:
|
||||
if (
|
||||
isinstance(tool_call, ChatCompletionMessageToolCall)
|
||||
and "persimmon" in tool_call.function.arguments
|
||||
):
|
||||
tool_call.function.arguments = '{"fruit": "[MASKED]"}'
|
||||
return inputs
|
||||
|
||||
|
|
@ -2602,9 +2684,7 @@ class TestAnthropicMessagesScanOnlyToolResults:
|
|||
{"role": "user", "content": "weather in Paris?"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{"type": "tool_use", "id": "tu1", "name": "get_weather", "input": {"city": "Paris"}}
|
||||
],
|
||||
"content": [{"type": "tool_use", "id": "tu1", "name": "get_weather", "input": {"city": "Paris"}}],
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
|
|
@ -2949,9 +3029,7 @@ class TestStructuredWriteBackKeepsToolResults:
|
|||
],
|
||||
ids=["tool_reference", "empty"],
|
||||
)
|
||||
async def test_tool_result_stays_right_after_its_tool_use(
|
||||
self, tool_result_content, expected_written_back_content
|
||||
):
|
||||
async def test_tool_result_stays_right_after_its_tool_use(self, tool_result_content, expected_written_back_content):
|
||||
handler = AnthropicMessagesHandler()
|
||||
data = {"model": "claude-fable-5", "messages": self._claude_code_tool_search_turns(tool_result_content)}
|
||||
|
||||
|
|
@ -3458,3 +3536,172 @@ def test_reading_a_dropped_budget_for_a_guardrail_logs_no_drop_warning(caplog: p
|
|||
|
||||
assert messages == [{"role": "user", "content": "hi"}]
|
||||
assert not [record for record in caplog.records if "Dropping" in record.getMessage()]
|
||||
|
||||
|
||||
def _tool_call_returning_guardrail(returned: Callable[[Sequence[object]], object]) -> CustomGuardrail:
|
||||
class ReturnToolCalls(CustomGuardrail):
|
||||
async def apply_guardrail(
|
||||
self,
|
||||
inputs: GenericGuardrailAPIInputs,
|
||||
request_data: dict[str, object],
|
||||
input_type: Literal["request", "response"],
|
||||
logging_obj: LiteLLMLoggingObj | None = None,
|
||||
) -> GenericGuardrailAPIInputs:
|
||||
return cast(
|
||||
GenericGuardrailAPIInputs,
|
||||
{
|
||||
**inputs,
|
||||
"texts": [f"[SCANNED] {text}" for text in inputs["texts"]],
|
||||
"tool_calls": returned(inputs["tool_calls"]),
|
||||
},
|
||||
)
|
||||
|
||||
return ReturnToolCalls(guardrail_name="tool-call-rewriter")
|
||||
|
||||
|
||||
def _with_arguments(tool_call: object, arguments: str) -> dict[str, object]:
|
||||
fields: Final = cast(dict[str, object], tool_call)
|
||||
function: Final = cast(dict[str, object], fields["function"])
|
||||
return {**fields, "function": {**function, "arguments": arguments}}
|
||||
|
||||
|
||||
def _redacting(tool_calls: Sequence[object]) -> list[dict[str, object]]:
|
||||
return [_with_arguments(tool_call, '{"cmd": "echo [REDACTED]"}') for tool_call in tool_calls]
|
||||
|
||||
|
||||
def _dict_tool_use_response() -> dict[str, object]:
|
||||
return {
|
||||
"id": "msg_1",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": "claude-sonnet-5-5",
|
||||
"content": [
|
||||
{"type": "text", "text": "Running it"},
|
||||
{"type": "tool_use", "id": "toolu_1", "name": "run_shell", "input": {"cmd": "echo SECRET-123"}},
|
||||
],
|
||||
"stop_reason": "tool_use",
|
||||
}
|
||||
|
||||
|
||||
def _object_tool_use_response() -> dict[str, object]:
|
||||
return {
|
||||
**_dict_tool_use_response(),
|
||||
"content": [
|
||||
AnthropicResponseContentBlockText(type="text", text="Running it"),
|
||||
AnthropicResponseContentBlockToolUse(
|
||||
type="tool_use", id="toolu_1", name="run_shell", input={"cmd": "echo SECRET-123"}
|
||||
),
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
def _tool_use_block(response: dict[str, object], index: int = 1) -> dict[str, object]:
|
||||
block: Final = cast(list[object], response["content"])[index]
|
||||
return cast(
|
||||
dict[str, object], block.model_dump() if isinstance(block, AnthropicResponseContentBlockToolUse) else block
|
||||
)
|
||||
|
||||
|
||||
class TestOutputToolCallRewrites:
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"response", [_dict_tool_use_response(), _object_tool_use_response()], ids=["dict_blocks", "object_blocks"]
|
||||
)
|
||||
async def test_rewritten_arguments_land_in_the_tool_use_input(self, response: dict[str, object]) -> None:
|
||||
await AnthropicMessagesHandler().process_output_response(
|
||||
response=cast(Any, response),
|
||||
guardrail_to_apply=_tool_call_returning_guardrail(_redacting),
|
||||
litellm_logging_obj=MagicMock(),
|
||||
)
|
||||
|
||||
block: Final = _tool_use_block(response)
|
||||
assert (block["id"], block["name"], block["input"]) == ("toolu_1", "run_shell", {"cmd": "echo [REDACTED]"})
|
||||
assert "SECRET-123" not in json.dumps(response, default=lambda block: block.model_dump())
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_rewritten_name_lands_on_the_tool_use_block(self) -> None:
|
||||
response: Final = _dict_tool_use_response()
|
||||
|
||||
def rename(tool_calls: Sequence[object]) -> list[dict[str, object]]:
|
||||
fields: Final = cast(dict[str, object], tool_calls[0])
|
||||
function: Final = cast(dict[str, object], fields["function"])
|
||||
return [{**fields, "function": {**function, "name": "run_safely"}}]
|
||||
|
||||
await AnthropicMessagesHandler().process_output_response(
|
||||
response=cast(Any, response),
|
||||
guardrail_to_apply=_tool_call_returning_guardrail(rename),
|
||||
litellm_logging_obj=MagicMock(),
|
||||
)
|
||||
|
||||
assert _tool_use_block(response)["name"] == "run_safely"
|
||||
assert _tool_use_block(response)["input"] == {"cmd": "echo SECRET-123"}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_each_returned_call_lands_on_the_tool_use_block_in_the_same_position(self) -> None:
|
||||
response: Final = _dict_tool_use_response()
|
||||
content: Final = cast(list[object], response["content"])
|
||||
content.append({"type": "text", "text": "and then"})
|
||||
content.append({"type": "tool_use", "id": "toolu_2", "name": "lookup", "input": {"query": "SECRET-123"}})
|
||||
|
||||
def redact_second_only(tool_calls: Sequence[object]) -> list[object]:
|
||||
return [tool_calls[0], _with_arguments(tool_calls[1], '{"query": "[REDACTED]"}')]
|
||||
|
||||
await AnthropicMessagesHandler().process_output_response(
|
||||
response=cast(Any, response),
|
||||
guardrail_to_apply=_tool_call_returning_guardrail(redact_second_only),
|
||||
litellm_logging_obj=MagicMock(),
|
||||
)
|
||||
|
||||
assert _tool_use_block(response, 1)["input"] == {"cmd": "echo SECRET-123"}
|
||||
assert _tool_use_block(response, 3)["input"] == {"query": "[REDACTED]"}
|
||||
assert [block["text"] for block in content if cast(dict[str, object], block)["type"] == "text"] == [
|
||||
"[SCANNED] Running it",
|
||||
"[SCANNED] and then",
|
||||
]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_rewritten_arguments_that_are_not_a_json_object_fail_the_response(self) -> None:
|
||||
response: Final = _dict_tool_use_response()
|
||||
|
||||
def break_arguments(tool_calls: Sequence[object]) -> list[dict[str, object]]:
|
||||
return [_with_arguments(tool_calls[0], "echo [REDACTED]")]
|
||||
|
||||
with pytest.raises(UnappliableRequestRewrite, match="tool-call-rewriter"):
|
||||
await AnthropicMessagesHandler().process_output_response(
|
||||
response=cast(Any, response),
|
||||
guardrail_to_apply=_tool_call_returning_guardrail(break_arguments),
|
||||
litellm_logging_obj=MagicMock(),
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"returned",
|
||||
[lambda tool_calls: [], lambda tool_calls: list(_redacting(tool_calls)) * 2, lambda tool_calls: None],
|
||||
ids=["empty", "doubled", "none"],
|
||||
)
|
||||
async def test_a_returned_list_that_cannot_be_placed_leaves_the_tool_use_input_as_scanned(
|
||||
self, returned: Callable[[Sequence[object]], object]
|
||||
) -> None:
|
||||
response: Final = _dict_tool_use_response()
|
||||
|
||||
await AnthropicMessagesHandler().process_output_response(
|
||||
response=cast(Any, response),
|
||||
guardrail_to_apply=_tool_call_returning_guardrail(returned),
|
||||
litellm_logging_obj=MagicMock(),
|
||||
)
|
||||
|
||||
assert _tool_use_block(response)["input"] == {"cmd": "echo SECRET-123"}
|
||||
assert _tool_use_block(response, 0) == {"type": "text", "text": "[SCANNED] Running it"}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_echoed_tool_calls_leave_the_tool_use_block_untouched(self) -> None:
|
||||
response: Final = _dict_tool_use_response()
|
||||
original_input: Final = _tool_use_block(response)["input"]
|
||||
|
||||
await AnthropicMessagesHandler().process_output_response(
|
||||
response=cast(Any, response),
|
||||
guardrail_to_apply=_tool_call_returning_guardrail(list),
|
||||
litellm_logging_obj=MagicMock(),
|
||||
)
|
||||
|
||||
assert _tool_use_block(response)["input"] is original_input
|
||||
|
|
|
|||
|
|
@ -2692,3 +2692,111 @@ async def test_during_mcp_call_blocked_raises() -> None:
|
|||
)
|
||||
|
||||
assert exc_info.value.message == "credential detected in tool input"
|
||||
|
||||
|
||||
def _run_shell_call(arguments: str) -> dict[str, object]:
|
||||
return {"id": "call_1", "type": "function", "function": {"name": "run_shell", "arguments": arguments}}
|
||||
|
||||
|
||||
def _guardrail_answering(body: dict[str, object], fail_on_error: bool = True) -> GenericGuardrailAPI:
|
||||
def answer(request: httpx.Request) -> httpx.Response:
|
||||
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,
|
||||
fail_on_error=fail_on_error,
|
||||
async_handler=AsyncHTTPHandler(transport=httpx.MockTransport(answer)),
|
||||
)
|
||||
|
||||
|
||||
async def _tool_calls_after(guardrail: GenericGuardrailAPI) -> GenericGuardrailAPIInputs:
|
||||
return await guardrail.apply_guardrail(
|
||||
inputs={"texts": [], "tool_calls": [_run_shell_call('{"cmd": "echo SECRET-123"}')]},
|
||||
request_data={},
|
||||
input_type="response",
|
||||
)
|
||||
|
||||
|
||||
class TestToolCallsInResponse:
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"returned_call",
|
||||
[
|
||||
_run_shell_call('{"cmd": "echo [REDACTED]"}'),
|
||||
{"function": {"name": "run_shell", "arguments": '{"cmd": "echo [REDACTED]"}'}},
|
||||
{
|
||||
"id": "ignored",
|
||||
"type": "ignored",
|
||||
"function": {"name": "run_shell", "arguments": '{"cmd": "echo [REDACTED]"}'},
|
||||
},
|
||||
],
|
||||
ids=["as_received", "function_only", "foreign_id_and_type"],
|
||||
)
|
||||
async def test_rewritten_tool_calls_reach_the_caller_with_the_calls_own_id_and_type(
|
||||
self, returned_call: dict[str, object]
|
||||
) -> None:
|
||||
guardrail: Final = _guardrail_answering({"action": "GUARDRAIL_INTERVENED", "tool_calls": [returned_call]})
|
||||
|
||||
guardrailed: Final = await _tool_calls_after(guardrail)
|
||||
|
||||
assert guardrailed["tool_calls"] == [_run_shell_call('{"cmd": "echo [REDACTED]"}')]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_rewritten_name_reaches_the_caller(self) -> None:
|
||||
guardrail: Final = _guardrail_answering(
|
||||
{"action": "GUARDRAIL_INTERVENED", "tool_calls": [{"function": {"name": "run_safely", "arguments": "{}"}}]}
|
||||
)
|
||||
|
||||
guardrailed: Final = await _tool_calls_after(guardrail)
|
||||
|
||||
assert guardrailed["tool_calls"] == [
|
||||
{"id": "call_1", "type": "function", "function": {"name": "run_safely", "arguments": "{}"}}
|
||||
]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"body",
|
||||
[
|
||||
{"action": "NONE"},
|
||||
{"action": "NONE", "tool_calls": None},
|
||||
{"action": "GUARDRAIL_INTERVENED", "texts": ["[MASKED]"], "tool_calls": []},
|
||||
],
|
||||
ids=["absent", "null", "empty"],
|
||||
)
|
||||
async def test_an_answer_without_tool_calls_leaves_the_key_out(self, body: dict[str, object]) -> None:
|
||||
guardrailed: Final = await _tool_calls_after(_guardrail_answering(body))
|
||||
|
||||
assert "tool_calls" not in guardrailed
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"tool_calls",
|
||||
[
|
||||
"not a list",
|
||||
["not a tool call"],
|
||||
[{"id": "call_1", "type": "function", "function": None}],
|
||||
[{"function": {"name": "run_shell"}}],
|
||||
[{"function": {"name": "run_shell", "arguments": {"cmd": "echo [REDACTED]"}}}],
|
||||
[_run_shell_call("{}"), _run_shell_call("{}")],
|
||||
],
|
||||
ids=["not_a_list", "not_a_call", "null_function", "no_arguments", "object_arguments", "one_call_too_many"],
|
||||
)
|
||||
async def test_a_tool_calls_answer_that_cannot_be_applied_fails_the_guardrail(self, tool_calls: object) -> None:
|
||||
guardrail: Final = _guardrail_answering({"action": "GUARDRAIL_INTERVENED", "tool_calls": tool_calls})
|
||||
|
||||
with pytest.raises(Exception, match="Generic Guardrail API failed"):
|
||||
await _tool_calls_after(guardrail)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fail_on_error_off_sends_the_calls_as_received_when_the_answer_cannot_be_applied(self) -> None:
|
||||
guardrail: Final = _guardrail_answering(
|
||||
{"action": "GUARDRAIL_INTERVENED", "tool_calls": [_run_shell_call("{}"), _run_shell_call("{}")]},
|
||||
fail_on_error=False,
|
||||
)
|
||||
|
||||
guardrailed: Final = await _tool_calls_after(guardrail)
|
||||
|
||||
assert guardrailed["tool_calls"] == [_run_shell_call('{"cmd": "echo SECRET-123"}')]
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue