mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-11 22:51:28 +00:00
fix(guardrails): fail open on Responses tool-call rewrites in another shape and keep item names
Guardrails that hand back tool_calls in their own shape (vendor JSON, user code output) raised a KeyError on the non-stream Responses write-back. Returned tool calls are now validated before comparison; a shape or count that does not line up leaves every tool-call item unchanged and logs a warning naming the guardrail. A tool call's name is written back only when the guardrail changed it, so a nameless custom_tool_call no longer picks up the custom_tool placeholder.
This commit is contained in:
mateo-berri
parent
ae6893bc64
commit
a248ff2c6e
2 changed files with 159 additions and 16 deletions
|
|
@ -37,7 +37,7 @@ from itertools import accumulate, chain, repeat
|
|||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Final, NamedTuple, Union, cast
|
||||
|
||||
from pydantic import BaseModel, TypeAdapter
|
||||
from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
|
||||
from typing_extensions import ReadOnly, TypedDict
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
|
|
@ -105,6 +105,19 @@ class _ToolCallShape(NamedTuple):
|
|||
arguments: str
|
||||
|
||||
|
||||
class _ToolCallFunctionFields(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
name: str | None = None
|
||||
arguments: str = ""
|
||||
|
||||
|
||||
class _ToolCallFields(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
function: _ToolCallFunctionFields
|
||||
|
||||
|
||||
def _tool_call_shapes(tool_calls: Sequence[ChatCompletionToolCallChunk]) -> tuple[_ToolCallShape, ...]:
|
||||
return tuple(
|
||||
_ToolCallShape(name=tool_call["function"].get("name"), arguments=tool_call["function"].get("arguments", ""))
|
||||
|
|
@ -112,6 +125,47 @@ def _tool_call_shapes(tool_calls: Sequence[ChatCompletionToolCallChunk]) -> tupl
|
|||
)
|
||||
|
||||
|
||||
def _returned_tool_call_shape(tool_call: object) -> _ToolCallShape | None:
|
||||
payload: Final = tool_call.model_dump() if isinstance(tool_call, BaseModel) else tool_call
|
||||
try:
|
||||
fields: Final = _ToolCallFields.model_validate(payload)
|
||||
except ValidationError:
|
||||
return None
|
||||
return _ToolCallShape(name=fields.function.name, arguments=fields.function.arguments)
|
||||
|
||||
|
||||
def _post_guardrail_tool_call_shapes(
|
||||
returned_tool_calls: Sequence[object] | None,
|
||||
pre_guardrail_tool_calls: tuple[_ToolCallShape, ...],
|
||||
guardrail_name: str | None,
|
||||
) -> tuple[_ToolCallShape, ...]:
|
||||
if not pre_guardrail_tool_calls:
|
||||
return pre_guardrail_tool_calls
|
||||
if returned_tool_calls is None or len(returned_tool_calls) != len(pre_guardrail_tool_calls):
|
||||
verbose_proxy_logger.warning(
|
||||
"OpenAI Responses API: guardrail %s returned %s tool calls for the %d scanned, "
|
||||
"leaving the tool call output items unchanged",
|
||||
guardrail_name,
|
||||
"no" if returned_tool_calls is None else len(returned_tool_calls),
|
||||
len(pre_guardrail_tool_calls),
|
||||
)
|
||||
return pre_guardrail_tool_calls
|
||||
returned_shapes: Final = tuple(_returned_tool_call_shape(tool_call) for tool_call in returned_tool_calls)
|
||||
validated_shapes: Final = tuple(shape for shape in returned_shapes if shape is not None)
|
||||
if len(validated_shapes) != len(returned_shapes):
|
||||
verbose_proxy_logger.warning(
|
||||
"OpenAI Responses API: guardrail %s returned tool calls without a function name and arguments, "
|
||||
"leaving the tool call output items unchanged",
|
||||
guardrail_name,
|
||||
)
|
||||
return pre_guardrail_tool_calls
|
||||
return validated_shapes
|
||||
|
||||
|
||||
def _tool_call_rewrite(before: _ToolCallShape, after: _ToolCallShape) -> _ToolCallShape:
|
||||
return _ToolCallShape(name=after.name if after.name != before.name else None, arguments=after.arguments)
|
||||
|
||||
|
||||
class ResponseOutputEnvelope(TypedDict, total=False):
|
||||
"""Dict form of a Responses API response, as far as guardrail write-back reads it."""
|
||||
|
||||
|
|
@ -698,7 +752,11 @@ class OpenAIResponsesHandler(BaseTranslation):
|
|||
)
|
||||
|
||||
guardrailed_texts: Final = guardrailed_inputs.get("texts", [])
|
||||
returned_tool_calls: Final = guardrailed_inputs.get("tool_calls")
|
||||
post_guardrail_tool_calls: Final = _post_guardrail_tool_call_shapes(
|
||||
returned_tool_calls=guardrailed_inputs.get("tool_calls"),
|
||||
pre_guardrail_tool_calls=pre_guardrail_tool_calls,
|
||||
guardrail_name=guardrail_to_apply.guardrail_name,
|
||||
)
|
||||
|
||||
# Step 3: Map guardrail responses back to original response structure
|
||||
await self._apply_guardrail_responses_to_output(
|
||||
|
|
@ -709,11 +767,7 @@ class OpenAIResponsesHandler(BaseTranslation):
|
|||
self._write_tool_call_rewrites_to_output(
|
||||
tool_call_items=tuple(item for item in response_output if _is_tool_call_output_item(item)),
|
||||
pre_guardrail_tool_calls=pre_guardrail_tool_calls,
|
||||
post_guardrail_tool_calls=_tool_call_shapes(
|
||||
returned_tool_calls
|
||||
if isinstance(returned_tool_calls, list) and len(returned_tool_calls) == len(tool_calls_to_check)
|
||||
else tool_calls_to_check
|
||||
),
|
||||
post_guardrail_tool_calls=post_guardrail_tool_calls,
|
||||
)
|
||||
|
||||
verbose_proxy_logger.debug("OpenAI Responses API: Processed output response: %s", response)
|
||||
|
|
@ -811,11 +865,10 @@ class OpenAIResponsesHandler(BaseTranslation):
|
|||
)
|
||||
|
||||
guardrailed_texts: Final = guardrailed_inputs.get("texts", [])
|
||||
returned_tool_calls: Final = guardrailed_inputs.get("tool_calls")
|
||||
post_guardrail_tool_calls: Final = _tool_call_shapes(
|
||||
returned_tool_calls
|
||||
if isinstance(returned_tool_calls, list) and len(returned_tool_calls) == len(tool_calls_to_check)
|
||||
else tool_calls_to_check
|
||||
post_guardrail_tool_calls: Final = _post_guardrail_tool_call_shapes(
|
||||
returned_tool_calls=guardrailed_inputs.get("tool_calls"),
|
||||
pre_guardrail_tool_calls=pre_guardrail_tool_calls,
|
||||
guardrail_name=guardrail_to_apply.guardrail_name,
|
||||
)
|
||||
|
||||
# Write guardrailed texts back into the output items in-place.
|
||||
|
|
@ -991,7 +1044,7 @@ class OpenAIResponsesHandler(BaseTranslation):
|
|||
)
|
||||
rewrites_by_call_id: Final = MappingProxyType(
|
||||
{
|
||||
call_id: after
|
||||
call_id: _tool_call_rewrite(before, after)
|
||||
for call_id, before, after in zip(call_ids, pre_guardrail_tool_calls, post_guardrail_tool_calls)
|
||||
if after != before
|
||||
}
|
||||
|
|
@ -1050,12 +1103,12 @@ class OpenAIResponsesHandler(BaseTranslation):
|
|||
) -> None:
|
||||
if len(tool_call_items) != len(post_guardrail_tool_calls):
|
||||
return
|
||||
for output_item, after in (
|
||||
(output_item, after)
|
||||
for output_item, rewrite in (
|
||||
(output_item, _tool_call_rewrite(before, after))
|
||||
for output_item, before, after in zip(tool_call_items, pre_guardrail_tool_calls, post_guardrail_tool_calls)
|
||||
if after != before
|
||||
):
|
||||
self._write_tool_call_item(output_item, after.name, after.arguments)
|
||||
self._write_tool_call_item(output_item, rewrite.name, rewrite.arguments)
|
||||
|
||||
@staticmethod
|
||||
def _tool_call_ids_by_item_id(stream_events: Sequence[object]) -> Mapping[str, str]:
|
||||
|
|
|
|||
|
|
@ -10,6 +10,8 @@ from collections.abc import Callable
|
|||
from typing import Any, List, Literal, Optional, Tuple
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import logging
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
|
|
@ -85,6 +87,29 @@ class PersimmonMaskingGuardrail(CustomGuardrail):
|
|||
return {**inputs, "tool_calls": tool_calls}
|
||||
|
||||
|
||||
class FlatShapeGuardrail(CustomGuardrail):
|
||||
async def apply_guardrail(
|
||||
self,
|
||||
inputs: GenericGuardrailAPIInputs,
|
||||
request_data: dict,
|
||||
input_type: Literal["request", "response"],
|
||||
logging_obj: Optional[LiteLLMLoggingObj] = None,
|
||||
) -> GenericGuardrailAPIInputs:
|
||||
flat_tool_calls = [{"name": "exec", "input": "rm -rf /"} for _ in inputs.get("tool_calls", [])]
|
||||
return {**inputs, "tool_calls": flat_tool_calls}
|
||||
|
||||
|
||||
class DroppingGuardrail(CustomGuardrail):
|
||||
async def apply_guardrail(
|
||||
self,
|
||||
inputs: GenericGuardrailAPIInputs,
|
||||
request_data: dict,
|
||||
input_type: Literal["request", "response"],
|
||||
logging_obj: Optional[LiteLLMLoggingObj] = None,
|
||||
) -> GenericGuardrailAPIInputs:
|
||||
return {**inputs, "tool_calls": []}
|
||||
|
||||
|
||||
CUSTOM_TOOL_CALL_ITEM = {
|
||||
"type": "custom_tool_call",
|
||||
"id": "ctc_1",
|
||||
|
|
@ -737,6 +762,52 @@ class TestOpenAIResponsesHandlerToolCallExtraction:
|
|||
assert (custom_item.name if typed else custom_item["name"]) == "exec"
|
||||
assert (output[0].content[0].text if typed else output[0]["content"][0]["text"]) == "running persimmon"
|
||||
|
||||
@staticmethod
|
||||
def _custom_tool_call_response(item: dict) -> dict:
|
||||
return {
|
||||
"id": "resp_1",
|
||||
"created_at": 1,
|
||||
"model": "gpt-5.6",
|
||||
"object": "response",
|
||||
"status": "completed",
|
||||
"output": [item],
|
||||
}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_process_output_response_ignores_tool_call_rewrites_in_another_shape(self):
|
||||
handler = OpenAIResponsesHandler()
|
||||
response = self._custom_tool_call_response(dict(CUSTOM_TOOL_CALL_ITEM))
|
||||
|
||||
result = await handler.process_output_response(response, FlatShapeGuardrail(guardrail_name="flat"))
|
||||
|
||||
assert result["output"][0]["input"] == "echo persimmon"
|
||||
assert result["output"][0]["name"] == "exec"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_process_output_response_warns_when_guardrail_drops_tool_calls(self, caplog):
|
||||
handler = OpenAIResponsesHandler()
|
||||
response = self._custom_tool_call_response(dict(CUSTOM_TOOL_CALL_ITEM))
|
||||
|
||||
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
|
||||
result = await handler.process_output_response(response, DroppingGuardrail(guardrail_name="dropper"))
|
||||
|
||||
assert result["output"][0]["input"] == "echo persimmon"
|
||||
assert any(
|
||||
"dropper" in record.getMessage() and "0 tool calls for the 1 scanned" in record.getMessage()
|
||||
for record in caplog.records
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_process_output_response_keeps_a_nameless_custom_tool_call_nameless(self):
|
||||
handler = OpenAIResponsesHandler()
|
||||
nameless_item = {key: value for key, value in CUSTOM_TOOL_CALL_ITEM.items() if key != "name"}
|
||||
response = self._custom_tool_call_response(nameless_item)
|
||||
|
||||
result = await handler.process_output_response(response, PersimmonMaskingGuardrail(guardrail_name="mask"))
|
||||
|
||||
assert result["output"][0]["input"] == "echo [MASKED]"
|
||||
assert "name" not in result["output"][0]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_process_output_response_with_tool_calls(self):
|
||||
"""Test processing output response containing function tool calls"""
|
||||
|
|
@ -1469,6 +1540,25 @@ class TestOpenAIResponsesHandlerStreamingOutputProcessing:
|
|||
assert events[5]["response"]["output"][0]["name"] == "exec"
|
||||
assert "arguments" not in events[5]["response"]["output"][0]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_deliver_ended_stream_rewrites_keep_a_nameless_custom_tool_call_nameless(self):
|
||||
handler = OpenAIResponsesHandler()
|
||||
events = self._ended_custom_tool_call_stream_events()
|
||||
items = [events[0]["item"], events[4]["item"], events[5]["response"]["output"][0]]
|
||||
for item in items:
|
||||
del item["name"]
|
||||
|
||||
await handler.process_output_streaming_response(
|
||||
responses_so_far=events,
|
||||
guardrail_to_apply=PersimmonMaskingGuardrail(guardrail_name="mask"),
|
||||
litellm_logging_obj=None,
|
||||
deliver_ended_stream_rewrites=True,
|
||||
)
|
||||
|
||||
assert events[3]["input"] == "echo [MASKED]"
|
||||
assert events[5]["response"]["output"][0]["input"] == "echo [MASKED]"
|
||||
assert all("name" not in item for item in items)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_deliver_ended_stream_rewrites_syncs_typed_custom_tool_call_events(self):
|
||||
from litellm.types.llms.openai import (
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue