Add image support

2026-10-03 02:22:24 +00:00 · 2026-04-13 14:11:42 -04:00 · 2026-04-13 14:11:42 -04:00 · ab6859245a
commit ab6859245a
parent c878dd51e2
2 changed files with 298 additions and 4 deletions
--- a/litellm/proxy/guardrails/guardrail_hooks/hiddenlayer/hiddenlayer.py
+++ b/litellm/proxy/guardrails/guardrail_hooks/hiddenlayer/hiddenlayer.py
@ -161,7 +161,7 @@ class HiddenlayerGuardrail(CustomGuardrail):
                    "messages": [
                        {
                            "role": last_msg.get("role", "user"),
-                            "content": last_msg.get("content", ""),
+                            "content": str(last_msg.get("content", "")),
                        }
                    ]
                },
@ -197,11 +197,29 @@ class HiddenlayerGuardrail(CustomGuardrail):
        if result.get("evaluation", {}).get("action") == HiddenlayerAction.REDACT:
            modified_data = result.get("modified_data", {})
            if modified_data.get("input") and input_type == "request":
-                inputs["texts"] = [modified_data["input"]["messages"][-1]["content"]]
+                last_content = modified_data["input"]["messages"][-1]["content"]
+                if isinstance(last_content, list):
+                    texts = [
+                        item["text"]
+                        for item in last_content
+                        if isinstance(item, dict) and item.get("type") == "text"
+                    ]
+                    inputs["texts"] = texts if texts else [""]
+                else:
+                    inputs["texts"] = [last_content]
                inputs["structured_messages"] = modified_data["input"]["messages"]

            if modified_data.get("output") and input_type == "response":
-                inputs["texts"] = [modified_data["output"]["messages"][-1]["content"]]
+                last_content = modified_data["output"]["messages"][-1]["content"]
+                if isinstance(last_content, list):
+                    texts = [
+                        item["text"]
+                        for item in last_content
+                        if isinstance(item, dict) and item.get("type") == "text"
+                    ]
+                    inputs["texts"] = texts if texts else [""]
+                else:
+                    inputs["texts"] = [last_content]

        return inputs

@ -414,7 +432,16 @@ class HiddenlayerGuardrailV2(CustomGuardrail):
            inputs["structured_messages"] = output

            for message in output.get("messages", []):
-                if content := message.get("content", ""):
+                content = message.get("content", "")
+                if isinstance(content, list):
+                    text_parts = [
+                        item["text"]
+                        for item in content
+                        if isinstance(item, dict) and item.get("type") == "text"
+                    ]
+                    if text_parts:
+                        new_texts.append(" ".join(text_parts))
+                elif content:
                    new_texts.append(content)

            inputs["texts"] = new_texts
--- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_hiddenlayer.py
+++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_hiddenlayer.py
@ -432,6 +432,138 @@ class TestHiddenlayerGuardrail:
                },
            )

+    @pytest.mark.asyncio
+    async def test_apply_guardrail_request_with_image(self):
+        """Test apply_guardrail sends multimodal content (image) to HiddenLayer v1."""
+        os.environ["HIDDENLAYER_API_BASE"] = "https://my.hiddenlayer"
+
+        guardrail = HiddenlayerGuardrail(
+            guardrail_name="hiddenlayer", event_hook="pre_call", default_on=True
+        )
+
+        multimodal_content = [
+            {"type": "text", "text": "how much is on this receipt?"},
+            {
+                "type": "image_url",
+                "image_url": {"url": "data:image/png;base64,iVBORw0KGgo="},
+            },
+        ]
+        inputs = GenericGuardrailAPIInputs(
+            texts=["how much is on this receipt?"],
+            images=["data:image/png;base64,iVBORw0KGgo="],
+            structured_messages=[{"role": "user", "content": multimodal_content}],
+            model="gpt-4o-mini",
+        )
+
+        request_data = {
+            "proxy_server_request": {
+                "headers": {},
+                "messages": [{"role": "user", "content": multimodal_content}],
+                "model": "gpt-4o-mini",
+            }
+        }
+
+        logging_obj = LiteLLMLoggingObj(
+            model="gpt-4o-mini",
+            messages=[{"role": "user", "content": multimodal_content}],
+            stream=False,
+            call_type="completion",
+            litellm_call_id="test-call-id",
+            function_id="test-function-id",
+            start_time=None,
+        )
+
+        mock_response = MagicMock()
+        mock_response.json.return_value = {}
+        mock_response.raise_for_status = MagicMock()
+
+        with patch.object(
+            guardrail._http_client, "post", return_value=mock_response
+        ) as mock_post:
+            result = await guardrail.apply_guardrail(
+                inputs=inputs,
+                request_data=request_data,
+                input_type="request",
+                logging_obj=logging_obj,
+            )
+
+        # v1 API requires string content — multimodal list is stringified
+        mock_post.assert_called_once()
+        call_kwargs = mock_post.call_args.kwargs
+        sent_content = call_kwargs["json"]["input"]["messages"][0]["content"]
+        assert isinstance(sent_content, str)
+        assert sent_content == str(multimodal_content)
+
+        # Result should be returned without error
+        assert result is not None
+
+    @pytest.mark.asyncio
+    async def test_apply_guardrail_redact_with_image_content(self):
+        """Test that REDACT action with multimodal content extracts text properly into inputs['texts']."""
+        os.environ["HIDDENLAYER_API_BASE"] = "https://my.hiddenlayer"
+
+        guardrail = HiddenlayerGuardrail(
+            guardrail_name="hiddenlayer", event_hook="pre_call", default_on=True
+        )
+
+        multimodal_content = [
+            {"type": "text", "text": "how much is on this receipt?"},
+            {
+                "type": "image_url",
+                "image_url": {"url": "data:image/png;base64,iVBORw0KGgo="},
+            },
+        ]
+        inputs = GenericGuardrailAPIInputs(
+            texts=["how much is on this receipt?"],
+            images=["data:image/png;base64,iVBORw0KGgo="],
+            structured_messages=[{"role": "user", "content": multimodal_content}],
+            model="gpt-4o-mini",
+        )
+
+        request_data = {"proxy_server_request": {"headers": {}}}
+
+        logging_obj = LiteLLMLoggingObj(
+            model="gpt-4o-mini",
+            messages=[],
+            stream=False,
+            call_type="completion",
+            litellm_call_id="test-call-id",
+            function_id="test-function-id",
+            start_time=None,
+        )
+
+        redacted_content = [
+            {"type": "text", "text": "[REDACTED]"},
+            {
+                "type": "image_url",
+                "image_url": {"url": "data:image/png;base64,iVBORw0KGgo="},
+            },
+        ]
+        mock_response = MagicMock()
+        mock_response.json.return_value = {
+            "evaluation": {"action": "Redact"},
+            "modified_data": {
+                "input": {
+                    "messages": [{"role": "user", "content": redacted_content}]
+                }
+            },
+        }
+        mock_response.raise_for_status = MagicMock()
+
+        with patch.object(guardrail._http_client, "post", return_value=mock_response):
+            result = await guardrail.apply_guardrail(
+                inputs=inputs,
+                request_data=request_data,
+                input_type="request",
+                logging_obj=logging_obj,
+            )
+
+        # texts must be List[str], not List[List]
+        assert result.get("texts") == ["[REDACTED]"]
+        assert result.get("structured_messages") == [
+            {"role": "user", "content": redacted_content}
+        ]
+
    def test_get_config_model(self):
        """Test get_config_model method."""
        config_model = HiddenlayerGuardrail.get_config_model()
@ -832,6 +964,141 @@ class TestHiddenlayerGuardrailV2:
                "detection/v2/response-evaluations" in mock_post.call_args.args[0]
            )

+    @pytest.mark.asyncio
+    async def test_apply_guardrail_request_with_image(self):
+        """Test apply_guardrail sends multimodal content (image) to HiddenLayer v2."""
+        os.environ["HIDDENLAYER_API_BASE"] = "https://my.hiddenlayer"
+
+        guardrail = HiddenlayerGuardrailV2(
+            guardrail_name="hiddenlayer", event_hook="pre_call", default_on=True
+        )
+
+        multimodal_content = [
+            {"type": "text", "text": "how much is on this receipt?"},
+            {
+                "type": "image_url",
+                "image_url": {"url": "data:image/png;base64,iVBORw0KGgo="},
+            },
+        ]
+        inputs = GenericGuardrailAPIInputs(
+            texts=["how much is on this receipt?"],
+            images=["data:image/png;base64,iVBORw0KGgo="],
+            structured_messages=[{"role": "user", "content": multimodal_content}],
+            model="gpt-4o-mini",
+        )
+
+        request_data = {
+            "proxy_server_request": {
+                "headers": {},
+                "messages": [{"role": "user", "content": multimodal_content}],
+                "model": "gpt-4o-mini",
+            }
+        }
+
+        logging_obj = LiteLLMLoggingObj(
+            model="gpt-4o-mini",
+            messages=[{"role": "user", "content": multimodal_content}],
+            stream=False,
+            call_type="completion",
+            litellm_call_id="test-call-id",
+            function_id="test-function-id",
+            start_time=None,
+        )
+
+        mock_response = MagicMock()
+        mock_response.headers = MagicMock()
+        mock_response.headers.get = MagicMock(return_value="")
+        mock_response.json.return_value = {
+            "messages": [{"role": "user", "content": multimodal_content}],
+            "model": "gpt-4o-mini",
+        }
+        mock_response.raise_for_status = MagicMock()
+
+        with patch.object(
+            guardrail._http_client, "post", return_value=mock_response
+        ) as mock_post:
+            result = await guardrail.apply_guardrail(
+                inputs=inputs,
+                request_data=request_data,
+                input_type="request",
+                logging_obj=logging_obj,
+            )
+
+        # Image data should be sent to HiddenLayer in the message content
+        mock_post.assert_called_once()
+        call_kwargs = mock_post.call_args.kwargs
+        sent_messages = call_kwargs["json"]["messages"]
+        assert sent_messages[0]["content"] == multimodal_content
+
+        # texts must be List[str] even when content is multimodal
+        texts = result.get("texts", [])
+        assert all(isinstance(t, str) for t in texts)
+        assert texts == ["how much is on this receipt?"]
+
+    @pytest.mark.asyncio
+    async def test_apply_guardrail_request_with_image_multimodal_response(self):
+        """Test that new_texts extraction handles multimodal content (list) returned by HiddenLayer v2."""
+        os.environ["HIDDENLAYER_API_BASE"] = "https://my.hiddenlayer"
+
+        guardrail = HiddenlayerGuardrailV2(
+            guardrail_name="hiddenlayer", event_hook="pre_call", default_on=True
+        )
+
+        multimodal_content = [
+            {"type": "text", "text": "how much is on this receipt?"},
+            {
+                "type": "image_url",
+                "image_url": {"url": "data:image/png;base64,iVBORw0KGgo="},
+            },
+        ]
+        inputs = GenericGuardrailAPIInputs(
+            texts=["how much is on this receipt?"],
+            images=["data:image/png;base64,iVBORw0KGgo="],
+            structured_messages=[{"role": "user", "content": multimodal_content}],
+            model="gpt-4o-mini",
+        )
+
+        request_data = {
+            "proxy_server_request": {
+                "headers": {},
+            }
+        }
+
+        logging_obj = LiteLLMLoggingObj(
+            model="gpt-4o-mini",
+            messages=[],
+            stream=False,
+            call_type="completion",
+            litellm_call_id="test-call-id",
+            function_id="test-function-id",
+            start_time=None,
+        )
+
+        # HiddenLayer returns the message with multimodal content unchanged
+        mock_response = MagicMock()
+        mock_response.headers = MagicMock()
+        mock_response.headers.get = MagicMock(return_value="")
+        mock_response.json.return_value = {
+            "messages": [{"role": "user", "content": multimodal_content}],
+            "model": "gpt-4o-mini",
+        }
+        mock_response.raise_for_status = MagicMock()
+
+        with patch.object(guardrail._http_client, "post", return_value=mock_response):
+            result = await guardrail.apply_guardrail(
+                inputs=inputs,
+                request_data=request_data,
+                input_type="request",
+                logging_obj=logging_obj,
+            )
+
+        # texts must be List[str], not List[List]
+        texts = result.get("texts", [])
+        assert all(isinstance(t, str) for t in texts), (
+            f"inputs['texts'] must be List[str], got: {texts}"
+        )
+        assert texts == ["how much is on this receipt?"]
+
    def test_get_config_model(self):
        """Test get_config_model method."""
        config_model = HiddenlayerGuardrailV2.get_config_model()