From 6bb63525dba41bf81d2c9e68f12653628efbd5fd Mon Sep 17 00:00:00 2001 From: Igal Boxerman Date: Mon, 12 Jan 2026 12:57:54 +0200 Subject: [PATCH] fix(guardrails): fix SerializationIterator error and pass tools to guardrail (#18932) * fix(generic-guardrail-api): fix SerializationIterator error on multimodal requests When sending multimodal messages (with images) through the Generic Guardrail API, the `model_dump()` call fails with "Object of type SerializationIterator is not JSON serializable" error. Root cause: The `ChatCompletionAssistantMessage` type defines `content` as an `Iterable` (not just `List`), and Pydantic's `model_dump()` creates a `SerializationIterator` for iterables which is not JSON serializable. Fix: Use `model_dump(mode="json")` which properly converts all iterables to lists and ensures all complex objects are JSON serializable. * fix(guardrails): pass tools (function definitions) to guardrail inputs The unified guardrail handler was not passing the `tools` parameter (function definitions) from the request to the guardrail inputs. This meant guardrails could not inspect or validate tool definitions. Added extraction of `data.get("tools")` and inclusion in the GenericGuardrailAPIInputs passed to `apply_guardrail()`. * test(guardrails): add tests for tools passed to guardrail Added tests verifying that tools (function definitions) are correctly passed to guardrails in the unified guardrail handler: - test_tools_passed_to_guardrail - test_multiple_tools_passed_to_guardrail - test_no_tools_in_request - test_tools_and_tool_calls_both_passed --- .../chat/guardrail_translation/handler.py | 4 + .../generic_guardrail_api.py | 3 +- test_generic_guardrail_config.yaml | 29 ++++ .../guardrail_translation/test_handler.py | 152 ++++++++++++++++++ .../test_generic_guardrail_api.py | 128 +++++++++++++++ 5 files changed, 315 insertions(+), 1 deletion(-) create mode 100644 test_generic_guardrail_config.yaml diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index 6c573894f6..e2cf9f4610 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -83,6 +83,10 @@ class OpenAIChatCompletionsHandler(BaseTranslation): inputs["structured_messages"] = ( messages # pass the openai /chat/completions messages to the guardrail, as-is ) + # Pass tools (function definitions) to the guardrail + tools = data.get("tools") + if tools: + inputs["tools"] = tools guardrailed_inputs = await guardrail_to_apply.apply_guardrail( inputs=inputs, 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 0dd00bfe55..41e0a23b9a 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 @@ -223,9 +223,10 @@ class GenericGuardrailAPI(CustomGuardrail): try: # Make the API request + # Use mode="json" to ensure all iterables are converted to lists response = await self.async_handler.post( url=self.api_base, - json=guardrail_request.model_dump(), + json=guardrail_request.model_dump(mode="json"), headers=headers, ) diff --git a/test_generic_guardrail_config.yaml b/test_generic_guardrail_config.yaml new file mode 100644 index 0000000000..d6cb505f7e --- /dev/null +++ b/test_generic_guardrail_config.yaml @@ -0,0 +1,29 @@ +model_list: + - model_name: gpt-4 + litellm_params: + model: openai/gpt-4 + api_key: os.environ/OPENAI_API_KEY + + - model_name: gpt-4o + litellm_params: + model: openai/gpt-4o + api_key: os.environ/OPENAI_API_KEY + + - model_name: gpt-3.5-turbo + litellm_params: + model: openai/gpt-3.5-turbo + api_key: os.environ/OPENAI_API_KEY + +guardrails: + - guardrail_name: thisispillar + litellm_params: + guardrail: generic_guardrail_api + mode: [pre_call, post_call] + api_base: os.environ/PILLAR_API_BASE + api_key: os.environ/PILLAR_API_KEY + default_on: true + additional_provider_specific_params: + plr_evidence: true + +general_settings: + master_key: sk-1234 diff --git a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_handler.py b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_handler.py index e94f40838c..6c0195d283 100644 --- a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_handler.py +++ b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_handler.py @@ -84,6 +84,158 @@ class MockGuardrail(CustomGuardrail): return result +class TestOpenAIChatCompletionsHandlerToolsInput: + """Test input processing with tools (function definitions)""" + + @pytest.mark.asyncio + async def test_tools_passed_to_guardrail(self): + """Test that tools (function definitions) are passed to the guardrail""" + handler = OpenAIChatCompletionsHandler() + guardrail = MockGuardrail() + + # Create input data with tools (function definitions) + data = { + "messages": [ + {"role": "user", "content": "What's the weather in Boston?"}, + ], + "tools": [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the current weather in a location", + "parameters": { + "type": "object", + "properties": { + "location": { + "type": "string", + "description": "City name", + }, + "unit": { + "type": "string", + "enum": ["celsius", "fahrenheit"], + }, + }, + "required": ["location"], + }, + }, + } + ], + } + + # Process the input + await handler.process_input_messages(data, guardrail) + + # Verify tools were passed to guardrail + assert guardrail.last_inputs is not None + assert "tools" in guardrail.last_inputs + assert len(guardrail.last_inputs["tools"]) == 1 + + tool = guardrail.last_inputs["tools"][0] + assert tool["type"] == "function" + assert tool["function"]["name"] == "get_weather" + assert tool["function"]["description"] == "Get the current weather in a location" + assert "parameters" in tool["function"] + + @pytest.mark.asyncio + async def test_multiple_tools_passed_to_guardrail(self): + """Test that multiple tools are passed to the guardrail""" + handler = OpenAIChatCompletionsHandler() + guardrail = MockGuardrail() + + data = { + "messages": [ + {"role": "user", "content": "What's the weather and time?"}, + ], + "tools": [ + { + "type": "function", + "function": { + "name": "get_weather", + "parameters": {"type": "object", "properties": {}}, + }, + }, + { + "type": "function", + "function": { + "name": "get_time", + "parameters": {"type": "object", "properties": {}}, + }, + }, + ], + } + + await handler.process_input_messages(data, guardrail) + + assert guardrail.last_inputs is not None + assert "tools" in guardrail.last_inputs + assert len(guardrail.last_inputs["tools"]) == 2 + assert guardrail.last_inputs["tools"][0]["function"]["name"] == "get_weather" + assert guardrail.last_inputs["tools"][1]["function"]["name"] == "get_time" + + @pytest.mark.asyncio + async def test_no_tools_in_request(self): + """Test that requests without tools work correctly""" + handler = OpenAIChatCompletionsHandler() + guardrail = MockGuardrail() + + data = { + "messages": [ + {"role": "user", "content": "Hello"}, + ], + } + + await handler.process_input_messages(data, guardrail) + + assert guardrail.last_inputs is not None + # tools should not be in inputs if not provided + assert "tools" not in guardrail.last_inputs or guardrail.last_inputs.get("tools") is None + + @pytest.mark.asyncio + async def test_tools_and_tool_calls_both_passed(self): + """Test that both tools (definitions) and tool_calls (invocations) are passed""" + handler = OpenAIChatCompletionsHandler() + guardrail = MockGuardrail() + + data = { + "messages": [ + {"role": "user", "content": "What's the weather?"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_123", + "type": "function", + "function": { + "name": "get_weather", + "arguments": '{"location": "Boston"}', + }, + } + ], + }, + ], + "tools": [ + { + "type": "function", + "function": { + "name": "get_weather", + "parameters": {"type": "object", "properties": {"location": {"type": "string"}}}, + }, + } + ], + } + + await handler.process_input_messages(data, guardrail) + + assert guardrail.last_inputs is not None + # Both should be present + assert "tools" in guardrail.last_inputs + assert "tool_calls" in guardrail.last_inputs + assert len(guardrail.last_inputs["tools"]) == 1 + assert len(guardrail.last_inputs["tool_calls"]) == 1 + + class TestOpenAIChatCompletionsHandlerToolCallsInput: """Test input processing with 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 b65065be36..8bf84b251d 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 @@ -589,3 +589,131 @@ class TestErrorHandling: ) assert "Generic Guardrail API failed" in str(exc_info.value) + + +class TestMultimodalSupport: + """Test multimodal (image) message handling and serialization""" + + @pytest.mark.asyncio + async def test_multimodal_message_serialization(self): + """ + Test that multimodal messages with images are properly serialized. + + This tests the fix for SerializationIterator error when messages contain + image_url content that includes Iterable types. + """ + guardrail = GenericGuardrailAPI( + api_base="https://api.test.guardrail.com", + guardrail_name="test-multimodal-guardrail", + ) + + # Create multimodal request data with image content + request_data = { + "model": "gpt-4o", + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "What's in this image?"}, + { + "type": "image_url", + "image_url": {"url": "https://example.com/image.jpg"}, + }, + ], + } + ], + "metadata": { + "user_api_key_user_id": "test-user", + }, + } + + mock_response = MagicMock() + mock_response.json.return_value = { + "action": "NONE", + "texts": ["What's in this image?"], + "images": ["https://example.com/image.jpg"], + } + mock_response.raise_for_status = MagicMock() + + with patch.object( + guardrail.async_handler, "post", return_value=mock_response + ) as mock_post: + # This should not raise SerializationIterator error + result = await guardrail.apply_guardrail( + inputs={ + "texts": ["What's in this image?"], + "images": ["https://example.com/image.jpg"], + "structured_messages": request_data["messages"], + }, + request_data=request_data, + input_type="request", + ) + + # Verify API was called successfully + mock_post.assert_called_once() + + # Verify the request was properly serialized (no SerializationIterator) + call_args = mock_post.call_args + json_payload = call_args.kwargs["json"] + + # Verify structured_messages is a proper list, not an iterator + assert isinstance(json_payload["structured_messages"], list) + assert json_payload["images"] == ["https://example.com/image.jpg"] + assert json_payload["texts"] == ["What's in this image?"] + + @pytest.mark.asyncio + async def test_iterable_content_serialization(self): + """ + Test that Iterable content types are properly converted to lists. + + The ChatCompletionAssistantMessage type allows content to be an Iterable, + which caused SerializationIterator errors before the fix. + """ + guardrail = GenericGuardrailAPI( + api_base="https://api.test.guardrail.com", + guardrail_name="test-iterable-guardrail", + ) + + # Simulate a message with content that could be an iterable + def content_generator(): + yield {"type": "text", "text": "Hello"} + yield {"type": "text", "text": "World"} + + # Create request with generator-based content (simulating Iterable type) + messages_with_iterable = [ + { + "role": "user", + "content": list(content_generator()), # Convert to list for test + } + ] + + request_data = { + "model": "gpt-4", + "messages": messages_with_iterable, + } + + mock_response = MagicMock() + mock_response.json.return_value = { + "action": "NONE", + "texts": ["Hello", "World"], + } + mock_response.raise_for_status = MagicMock() + + with patch.object( + guardrail.async_handler, "post", return_value=mock_response + ) as mock_post: + result = await guardrail.apply_guardrail( + inputs={ + "texts": ["Hello", "World"], + "structured_messages": messages_with_iterable, + }, + request_data=request_data, + input_type="request", + ) + + mock_post.assert_called_once() + + # Verify serialization succeeded + call_args = mock_post.call_args + json_payload = call_args.kwargs["json"] + assert isinstance(json_payload["structured_messages"], list) \ No newline at end of file