Merge e5cab21b18 into eb9355d11a

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:

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")),
         )

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"""