fix(generic_guardrail_api): return tool_calls on the guardrail response

litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py

@@ -29,6 +29,7 @@
     GenericGuardrailAPIMetadata,
     GenericGuardrailAPIRequest,
     GenericGuardrailAPIResponse,
+    GenericGuardrailAPIToolCalls,
     GuardrailToolParam,
 )
 from litellm.types.utils import GenericGuardrailAPIInputs
@@ -104,6 +105,16 @@ def _sanitize_inbound_headers(
     return sanitized or None
 
 
+def _returned_tool_call_is_usable(tool_call: object) -> bool:
+    """Whether a guardrail's returned tool call carries what the chat handlers index."""
+    function: Final = tool_call.get("function") if isinstance(tool_call, dict) else getattr(tool_call, "function", None)
+    if function is None:
+        return False
+    name: Final = function.get("name") if isinstance(function, dict) else getattr(function, "name", None)
+    arguments: Final = function.get("arguments") if isinstance(function, dict) else getattr(function, "arguments", None)
+    return isinstance(name, str) and isinstance(arguments, str)
+
+
 def _extract_inbound_headers(
     request_data: dict,
     logging_obj: Optional["LiteLLMLoggingObj"],
@@ -323,6 +334,7 @@ def _build_guardrail_return_inputs(
         images: list[str] | None,
         tools: list[ChatCompletionToolParam] | None,
         guardrail_response: GenericGuardrailAPIResponse,
+        tool_calls: GenericGuardrailAPIToolCalls | None = None,
     ) -> GenericGuardrailAPIInputs:
         # Action is NONE or no modifications needed
         return_inputs: Final = GenericGuardrailAPIInputs(texts=texts)
@@ -336,6 +348,16 @@ def _build_guardrail_return_inputs(
             return_inputs["tools"] = guardrail_response.tools
         elif tools:
             return_inputs["tools"] = tools
+        returned_tool_calls: Final = guardrail_response.tool_calls
+        if returned_tool_calls and not all(_returned_tool_call_is_usable(call) for call in returned_tool_calls):
+            verbose_proxy_logger.warning(
+                "Generic Guardrail API returned tool calls missing a function name or arguments; keeping the ones already sent",
+            )
+            returned_tool_calls = None  # rebind-ok: an unusable reply is discarded
+        if returned_tool_calls:
+            return_inputs["tool_calls"] = returned_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 +495,7 @@ async def apply_guardrail(
                 texts=texts,
                 images=images,
                 tools=tools,
+                tool_calls=tool_calls,
                 guardrail_response=guardrail_response,
             )
 

litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py

@@ -1,4 +1,4 @@
-from typing import Any, Final, Literal
+from typing import Any, Final, Literal, TypeAlias
 
 from pydantic import BaseModel, ConfigDict, Field
 from typing_extensions import TypedDict
@@ -8,7 +8,9 @@
     ChatCompletionToolCallChunk,
 )
 from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel
-from litellm.types.utils import ChatCompletionMessageToolCall
+from litellm.types.utils import ChatCompletionMessageToolCall, GenericGuardrailAPIInputs
+
+GenericGuardrailAPIToolCalls: TypeAlias = GenericGuardrailAPIInputs["tool_calls"]  # noqa: F821  # TypedDict field reference, not a name
 
 
 class GuardrailToolParam(BaseModel):
@@ -167,6 +169,7 @@ class GenericGuardrailAPIResponse:
     action: str
     blocked_reason: str | None
     stream_holdback_chars: list[int] | None
+    tool_calls: GenericGuardrailAPIToolCalls | None
 
     def __init__(
         self,
@@ -176,6 +179,7 @@ def __init__(
         images: list[str] | None = None,
         tools: list[GuardrailToolParam] | None = None,
         stream_holdback_chars: list[int] | None = None,
+        tool_calls: GenericGuardrailAPIToolCalls | None = None,
     ) -> None:
         self.action = action
         self.blocked_reason = blocked_reason
@@ -186,6 +190,7 @@ def __init__(
         # 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":
@@ -200,4 +205,5 @@ def from_dict(cls, data: dict) -> "GenericGuardrailAPIResponse":
             images=data.get("images"),
             tools=data.get("tools"),
             stream_holdback_chars=stream_holdback_chars,
+            tool_calls=data.get("tool_calls"),
         )

tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py

@@ -1380,6 +1380,35 @@ def test_from_dict_malformed_holdback_degrades_to_zero(self):
 
         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,141 @@ async def test_builtin_tools_without_function_block_do_not_crash(
 
         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_malformed_returned_tool_calls_are_discarded(self, generic_guardrail):
+        """A same-length list with no function block must not reach chat handling.
+
+        The chat handlers index ``function.name`` and ``function.arguments`` on whatever the
+        guardrail returns, so forwarding an unusable list raises there instead of returning a
+        response. The calls that were sent are kept instead.
+        """
+        from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import (
+            GenericGuardrailAPIResponse,
+        )
+
+        sent = [{"id": "call_1", "type": "function", "function": {"name": "send_email", "arguments": "{}"}}]
+        malformed = [{"id": "call_1", "type": "function"}]
+
+        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": malformed}
+            ),
+        )
+
+        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"""