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 d8296003ae9..8114ea9aa48 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 @@ -24,14 +24,14 @@ from litellm.llms.custom_httpx.http_handler import ( httpxSpecialProvider, ) from litellm.types.guardrails import GuardrailEventHooks -from litellm.types.llms.openai import ChatCompletionToolParam +from litellm.types.llms.openai import ChatCompletionToolCallChunk, ChatCompletionToolParam from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import ( GenericGuardrailAPIMetadata, GenericGuardrailAPIRequest, GenericGuardrailAPIResponse, GuardrailToolParam, ) -from litellm.types.utils import GenericGuardrailAPIInputs +from litellm.types.utils import ChatCompletionMessageToolCall, GenericGuardrailAPIInputs if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -323,6 +323,7 @@ class GenericGuardrailAPI(CustomGuardrail): images: list[str] | None, tools: list[ChatCompletionToolParam] | None, guardrail_response: GenericGuardrailAPIResponse, + tool_calls: list[ChatCompletionToolCallChunk] | list[ChatCompletionMessageToolCall] | None = None, ) -> GenericGuardrailAPIInputs: # Action is NONE or no modifications needed return_inputs: Final = GenericGuardrailAPIInputs(texts=texts) @@ -336,6 +337,14 @@ class GenericGuardrailAPI(CustomGuardrail): return_inputs["tools"] = guardrail_response.tools elif tools: return_inputs["tools"] = tools + # Mirrors ``tools``: the guardrail's replacements win, otherwise the calls + # are echoed back unchanged. Handlers compare the returned list against the + # one they sent and fall back to their own copy on a length mismatch, so a + # guardrail that cannot reorder them still cannot change their count. + if guardrail_response.tool_calls: + return_inputs["tool_calls"] = guardrail_response.tool_calls + elif tool_calls: + return_inputs["tool_calls"] = tool_calls if guardrail_response.stream_holdback_chars is not None: return_inputs["stream_holdback_chars"] = guardrail_response.stream_holdback_chars return return_inputs @@ -473,6 +482,7 @@ class GenericGuardrailAPI(CustomGuardrail): texts=texts, images=images, tools=tools, + tool_calls=tool_calls, guardrail_response=guardrail_response, ) 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 4a868c48352..350e4976984 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py @@ -167,6 +167,7 @@ class GenericGuardrailAPIResponse: action: str blocked_reason: str | None stream_holdback_chars: list[int] | None + tool_calls: list[ChatCompletionToolCallChunk] | list[ChatCompletionMessageToolCall] | None def __init__( self, @@ -176,6 +177,7 @@ class GenericGuardrailAPIResponse: images: list[str] | None = None, tools: list[GuardrailToolParam] | None = None, stream_holdback_chars: list[int] | None = None, + tool_calls: list[ChatCompletionToolCallChunk] | list[ChatCompletionMessageToolCall] | None = None, ) -> None: self.action = action self.blocked_reason = blocked_reason @@ -186,6 +188,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 + # Replacement tool calls. ``tool_calls`` is sent on the request, so the + # response mirrors it back the way ``tools`` and ``images`` are: without + # this field a guardrail that rewrites tool arguments has no channel to + # return them on, and the caller keeps the values it sent in. + self.tool_calls = tool_calls @classmethod def from_dict(cls, data: dict) -> "GenericGuardrailAPIResponse": @@ -200,4 +207,5 @@ class GenericGuardrailAPIResponse: images=data.get("images"), tools=data.get("tools"), stream_holdback_chars=stream_holdback_chars, + tool_calls=data.get("tool_calls"), ) diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py index 83cc9ae8bb9..df672569adb 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py @@ -1380,6 +1380,35 @@ class TestGenericGuardrailAPIResponseParsing: assert response.stream_holdback_chars == [3, 0, 0, 0] + def test_from_dict_parses_tool_calls(self): + """``tool_calls`` goes out on the request, so the response has to carry it back.""" + 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 @@ -1876,6 +1905,115 @@ class TestToolSupport: assert forwarded_tools == tools + @pytest.mark.asyncio + async def test_guardrail_returned_tool_calls_flow_back_to_inputs(self, generic_guardrail): + """A guardrail that rewrites tool arguments has to be able to return them. + + ``tool_calls`` is sent on the request, so the matching response field is what closes + the loop. Without it the calls were dropped on the floor and the caller silently + kept the values it sent in -- for a rewriting guardrail (reversible redaction, for + instance) the rewritten arguments never reached the client. + """ + 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): + """The returned calls win over the ones that went out. + + Built through ``_build_guardrail_return_inputs`` rather than ``apply_guardrail`` so + the assertion is about this mapping, not about request-model coercion of the calls + the caller supplied. + """ + 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, + tool_calls=sent, + guardrail_response=GenericGuardrailAPIResponse.from_dict( + {"action": "GUARDRAIL_INTERVENED", "texts": ["Done"], "tool_calls": returned} + ), + ) + + assert result["tool_calls"] == returned + + def test_untouched_tool_calls_are_echoed_not_dropped(self, generic_guardrail): + """A guardrail that ignores tool calls must not blank them out. + + The downstream handlers compare what comes back against what they sent and fall back + to their own copy when the lengths differ, so echoing an unchanged list leaves them + behaving exactly as they did before this field existed. + """ + from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import ( + GenericGuardrailAPIResponse, + ) + + sent = [ + { + "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, + tool_calls=sent, + guardrail_response=GenericGuardrailAPIResponse.from_dict({"action": "NONE", "texts": ["Done"]}), + ) + + assert result["tool_calls"] == sent + + def test_no_tool_calls_on_either_side_stay_absent(self, generic_guardrail): + """An empty list is not the same as no list: the key must stay out of the result.""" + 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, + 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"""