diff --git a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py index 7bb41b7586b..53a2ae4369d 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py +++ b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py @@ -336,6 +336,15 @@ class GenericGuardrailAPI(CustomGuardrail): headers.update(self.headers) return headers + @staticmethod + def _echo_sent_tool_calls( + guardrailed: GenericGuardrailAPIInputs, *, sent: GenericGuardrailAPIInputs + ) -> GenericGuardrailAPIInputs: + sent_tool_calls: Final = sent.get("tool_calls") + if not sent_tool_calls or "tool_calls" in guardrailed: + return guardrailed + return {**guardrailed, "tool_calls": sent_tool_calls} + def _build_guardrail_return_inputs( self, *, @@ -365,6 +374,8 @@ class GenericGuardrailAPI(CustomGuardrail): ) if rows_to_write_back is not None: return_inputs["structured_messages"] = list(rows_to_write_back) + if guardrail_response.tool_calls: + return_inputs["tool_calls"] = list(guardrail_response.tool_calls) if guardrail_response.stream_holdback_chars is not None: return_inputs["stream_holdback_chars"] = guardrail_response.stream_holdback_chars return return_inputs @@ -499,13 +510,16 @@ class GenericGuardrailAPI(CustomGuardrail): blocked_content=True, ) - return self._build_guardrail_return_inputs( - texts=texts, - images=images, - tools=tools, - structured_messages=structured_messages, - shown_messages=guardrail_request.structured_messages, - guardrail_response=guardrail_response, + return self._echo_sent_tool_calls( + self._build_guardrail_return_inputs( + texts=texts, + images=images, + tools=tools, + structured_messages=structured_messages, + shown_messages=guardrail_request.structured_messages, + guardrail_response=guardrail_response, + ), + sent=inputs, ) except GuardrailRaisedException: diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py b/litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py index 49eaa154fec..e055bf7f8a3 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py @@ -4,6 +4,7 @@ from typing import Any, Final, Literal, cast # noqa: TID251 # JSON chat rows h from pydantic import ConfigDict, Field from typing_extensions import TypedDict +from litellm._logging import verbose_proxy_logger from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.llms.openai import ( AllMessageValues, @@ -168,6 +169,40 @@ 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 +def _json_object(value: object) -> Mapping[str, object] | None: + if not isinstance(value, Mapping): + return None + return cast("Mapping[str, object]", value) # cast-ok: isinstance leaves a Mapping's item types unknown + + +def _json_field(container: object, key: str) -> object: + fields: Final = _json_object(container) + return fields.get(key) if fields is not None else None + + +def _is_usable_tool_call(row: object) -> bool: + function: Final = _json_field(row, "function") + return isinstance(_json_field(function, "name"), str) and isinstance(_json_field(function, "arguments"), str) + + +def _json_array(value: object) -> Sequence[object] | None: + if not isinstance(value, list): + return None + return cast("Sequence[object]", value) # cast-ok: isinstance leaves a list's item type unknown + + +def tool_calls_from_response(value: object) -> Sequence[ChatCompletionToolCallChunk] | None: + rows: Final = _json_array(value) + if rows is None: + return None + if not all(_is_usable_tool_call(row) for row in rows): + verbose_proxy_logger.warning( + "Generic Guardrail API returned tool calls missing a function name or arguments; keeping the ones already sent", + ) + return None + return cast("Sequence[ChatCompletionToolCallChunk]", rows) # cast-ok: every row checked above + + class GenericGuardrailAPIResponse: """Response model for the Generic Guardrail API""" @@ -178,6 +213,7 @@ class GenericGuardrailAPIResponse: action: str blocked_reason: str | None stream_holdback_chars: list[int] | None + tool_calls: Sequence[ChatCompletionToolCallChunk] | None def __init__( self, @@ -188,6 +224,7 @@ 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 @@ -199,9 +236,11 @@ class GenericGuardrailAPIResponse: # framework must withhold from streaming emission until the next # processing round (word-boundary safety for text transformations). self.stream_holdback_chars = stream_holdback_chars + self.tool_calls = tool_calls @classmethod def from_dict(cls, data: dict) -> "GenericGuardrailAPIResponse": + fields: Final = _json_object(data) or {} raw_holdback: Final = data.get("stream_holdback_chars") stream_holdback_chars: Final = ( [coerce_stream_holdback_value(value) for value in raw_holdback] if isinstance(raw_holdback, list) else None @@ -213,5 +252,6 @@ class GenericGuardrailAPIResponse: images=data.get("images"), tools=data.get("tools"), stream_holdback_chars=stream_holdback_chars, - structured_messages=structured_messages_from_response(data.get("structured_messages")), + structured_messages=structured_messages_from_response(fields.get("structured_messages")), + tool_calls=tool_calls_from_response(fields.get("tool_calls")), ) diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py index e97de4686bf..8f274b89f98 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py @@ -1519,6 +1519,34 @@ class TestGenericGuardrailAPIResponseParsing: assert response.stream_holdback_chars == [3, 0, 0, 0] + def test_from_dict_parses_tool_calls(self): + from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import ( + GenericGuardrailAPIResponse, + ) + + tool_calls = [ + { + "id": "call_1", + "type": "function", + "function": {"name": "send_email", "arguments": '{"to": "jane@example.com"}'}, + } + ] + + response = GenericGuardrailAPIResponse.from_dict( + {"action": "GUARDRAIL_INTERVENED", "texts": ["Done"], "tool_calls": tool_calls} + ) + + assert response.tool_calls == tool_calls + + def test_from_dict_tool_calls_absent_is_none(self): + from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import ( + GenericGuardrailAPIResponse, + ) + + response = GenericGuardrailAPIResponse.from_dict({"action": "NONE", "texts": ["hi"]}) + + assert response.tool_calls is None + @pytest.mark.asyncio async def test_apply_guardrail_flows_holdback_back_to_inputs(self, generic_guardrail): """A GUARDRAIL_INTERVENED response with stream_holdback_chars is surfaced on @@ -2015,6 +2043,132 @@ class TestToolSupport: assert forwarded_tools == tools + @pytest.mark.asyncio + async def test_guardrail_returned_tool_calls_flow_back_to_inputs(self, generic_guardrail): + guardrailed_tool_calls = [ + { + "id": "call_1", + "type": "function", + "function": {"name": "send_email", "arguments": '{"to": "jane@example.com"}'}, + } + ] + mock_response = MagicMock() + mock_response.json.return_value = { + "action": "GUARDRAIL_INTERVENED", + "texts": ["Done"], + "tool_calls": guardrailed_tool_calls, + } + mock_response.raise_for_status = MagicMock() + + with patch.object(generic_guardrail.async_handler, "post", return_value=mock_response): + result = await generic_guardrail.apply_guardrail( + inputs={"texts": ["Done"]}, + request_data={}, + input_type="response", + ) + + assert result["tool_calls"] == guardrailed_tool_calls + + def test_guardrail_tool_calls_take_precedence_over_the_sent_ones(self, generic_guardrail): + from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import ( + GenericGuardrailAPIResponse, + ) + + sent = [{"id": "call_1", "type": "function", "function": {"name": "send_email", "arguments": "{}"}}] + returned = [ + { + "id": "call_1", + "type": "function", + "function": {"name": "send_email", "arguments": '{"to": "jane@example.com"}'}, + } + ] + + result = generic_guardrail._build_guardrail_return_inputs( + texts=["Done"], + images=None, + tools=None, + structured_messages=None, + shown_messages=None, + guardrail_response=GenericGuardrailAPIResponse.from_dict( + {"action": "GUARDRAIL_INTERVENED", "texts": ["Done"], "tool_calls": returned} + ), + ) + + assert result["tool_calls"] == returned + assert result["tool_calls"] != sent + + @pytest.mark.asyncio + async def test_untouched_tool_calls_are_echoed_not_dropped(self, generic_guardrail): + sent = [ + { + "id": "call_1", + "type": "function", + "function": {"name": "send_email", "arguments": '{"to": "jane@example.com"}'}, + } + ] + mock_response = MagicMock() + mock_response.json.return_value = {"action": "NONE", "texts": ["Done"]} + mock_response.raise_for_status = MagicMock() + + with patch.object(generic_guardrail.async_handler, "post", return_value=mock_response): + result = await generic_guardrail.apply_guardrail( + inputs={"texts": ["Done"], "tool_calls": sent}, + request_data={}, + input_type="response", + ) + + assert result["tool_calls"] is sent + + @pytest.mark.asyncio + async def test_malformed_returned_tool_calls_are_discarded(self, generic_guardrail): + sent = [{"id": "call_1", "type": "function", "function": {"name": "send_email", "arguments": "{}"}}] + mock_response = MagicMock() + mock_response.json.return_value = { + "action": "GUARDRAIL_INTERVENED", + "texts": ["Done"], + "tool_calls": [{"id": "call_1", "type": "function"}], + } + mock_response.raise_for_status = MagicMock() + + with patch.object(generic_guardrail.async_handler, "post", return_value=mock_response): + result = await generic_guardrail.apply_guardrail( + inputs={"texts": ["Done"], "tool_calls": sent}, + request_data={}, + input_type="response", + ) + + assert result["tool_calls"] is sent + + def test_from_dict_rejects_tool_calls_without_string_arguments(self): + from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import ( + GenericGuardrailAPIResponse, + ) + + response = GenericGuardrailAPIResponse.from_dict( + { + "action": "GUARDRAIL_INTERVENED", + "tool_calls": [{"id": "call_1", "function": {"name": "send_email", "arguments": {"to": "x"}}}], + } + ) + + assert response.tool_calls is None + + def test_no_tool_calls_on_either_side_stay_absent(self, generic_guardrail): + from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import ( + GenericGuardrailAPIResponse, + ) + + result = generic_guardrail._build_guardrail_return_inputs( + texts=["Done"], + images=None, + tools=None, + structured_messages=None, + shown_messages=None, + guardrail_response=GenericGuardrailAPIResponse.from_dict({"action": "NONE", "texts": ["Done"]}), + ) + + assert "tool_calls" not in result + class TestFailOnError: """Test fail_on_error: complete fail-open on any guardrail error"""