Merge 058fd590f7 into 231a46e40b

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

@@ -1,7 +1,7 @@
-from collections.abc import Mapping, Sequence
+from collections.abc import Callable, Mapping, Sequence
 from typing import Any, Final, Literal, cast  # noqa: TID251  # JSON chat rows have no typed constructor across roles
 
-from pydantic import BaseModel, ConfigDict, Field
+from pydantic import BaseModel, ConfigDict, Field, model_serializer
 from typing_extensions import TypedDict
 
 from litellm.types.llms.openai import (
@@ -15,13 +15,23 @@ from litellm.types.utils import ChatCompletionMessageToolCall
 class GuardrailToolParam(BaseModel):
     """A tool forwarded verbatim to the guardrail for inspection.
 
-    Built-in tools (code_interpreter, file_search, ...) have no ``function`` block
-    and stash their config in tool-specific keys, so only ``type`` is required and
-    ``extra="allow"`` preserves the rest instead of stripping it.
+    ``extra="allow"`` keeps provider-specific keys. ``type`` is optional for
+    provider-native tools such as Gemini ``{"googleSearch": {}}``.
     """
 
     model_config = ConfigDict(extra="allow")
-    type: str
+    type: str | None = None
+
+    @model_serializer(mode="wrap")
+    def _omit_null_type(  # noqa: ANN202  # annotating it replaces the model's serialization schema
+        self, handler: Callable[[object], Mapping[str, object]]
+    ):
+        data: Final[Mapping[str, object]] = handler(self)
+        if not isinstance(data, dict) or data.get("type") is not None:
+            return data
+        return {  # mutable-ok: pydantic's json serializer rejects a mapping that is not a dict
+            key: value for key, value in data.items() if key != "type"
+        }
 
 
 class GenericGuardrailAPIMetadata(TypedDict, total=False):

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

@@ -2015,6 +2015,34 @@ class TestToolSupport:
 
         assert forwarded_tools == tools
 
+    @pytest.mark.asyncio
+    async def test_gemini_native_tools_without_type_do_not_crash(
+        self, generic_guardrail
+    ):
+        """Gemini-native tools without type are forwarded unchanged."""
+        tools = [
+            {"googleSearch": {}},
+            {"codeExecution": {}},
+            {"type": "function", "function": {"name": "get_weather", "parameters": {}}},
+        ]
+
+        mock_response = MagicMock()
+        mock_response.json.return_value = {"action": "NONE", "texts": ["hi"]}
+        mock_response.raise_for_status = MagicMock()
+
+        with patch.object(
+            generic_guardrail.async_handler, "post", return_value=mock_response
+        ) as mock_post:
+            await generic_guardrail.apply_guardrail(
+                inputs={"texts": ["hi"], "tools": tools},
+                request_data={},
+                input_type="request",
+            )
+
+            forwarded_tools = mock_post.call_args.kwargs["json"]["tools"]
+
+        assert forwarded_tools == tools
+
 
 class TestFailOnError:
     """Test fail_on_error: complete fail-open on any guardrail error"""