From ab6859245a8c4f3211291a2e7372b6af7b9d1fe1 Mon Sep 17 00:00:00 2001 From: Ashton Sidhu Date: Mon, 13 Apr 2026 14:11:42 -0400 Subject: [PATCH] Add image support --- .../hiddenlayer/hiddenlayer.py | 35 ++- .../guardrail_hooks/test_hiddenlayer.py | 267 ++++++++++++++++++ 2 files changed, 298 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/hiddenlayer/hiddenlayer.py b/litellm/proxy/guardrails/guardrail_hooks/hiddenlayer/hiddenlayer.py index 274fed7207b..9ea93fa667b 100644 --- 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 diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_hiddenlayer.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_hiddenlayer.py index 3ec07c12284..23cbf1c03b0 100644 --- 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()