From 797306ccb0d5eb4b50bc73c85926790d12ce37d6 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 31 Jul 2026 07:22:35 +0000 Subject: [PATCH] fix(guardrails): send image parts to bedrock ApplyGuardrail Image content items were dropped when building the ApplyGuardrail INPUT payload, so a guardrail with the IMAGE modality enabled only ever saw the text next to the image. Inline and remote images are now decoded through BedrockImageProcessor and sent as image blocks; formats ApplyGuardrail cannot take are skipped with a warning Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../guardrail_hooks/bedrock_guardrails.py | 118 +++++++++++++++--- .../guardrail_hooks/bedrock_guardrails.py | 13 ++ .../test_bedrock_guardrails.py | 4 +- .../test_bedrock_guardrails.py | 117 ++++++++++++++++- .../guardrails/test_guardrail_endpoints.py | 4 +- 5 files changed, 233 insertions(+), 23 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index cec682d772a..0bd496fd4b9 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -26,6 +26,7 @@ from typing import ( cast, ) +import asyncio import copy from collections.abc import Mapping from datetime import datetime, timezone @@ -40,6 +41,7 @@ from litellm.litellm_core_utils.core_helpers import redact_nested_match_and_rege from litellm.caching import DualCache from litellm.exceptions import ModifyResponseException from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.litellm_core_utils.prompt_templates.factory import BedrockImageProcessor from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, @@ -54,9 +56,12 @@ from litellm.types.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( BedrockChecksViolation, BedrockContentItem, BedrockGuardrailChecksResponse, + BedrockGuardrailImageFormat, + BedrockGuardrailImageSource, BedrockGuardrailOutput, BedrockGuardrailQualifier, BedrockGuardrailResponse, + BedrockImageContent, BedrockRequest, BedrockTextContent, ) @@ -115,6 +120,12 @@ _CONTENT_TYPE_TO_QUALIFIER: Dict[str, BedrockGuardrailQualifier] = { # must not be graded against as if it were the application's own source material. _GROUNDING_SOURCE_TRUSTED_ROLES = frozenset({"system", "developer"}) +_APPLY_GUARDRAIL_IMAGE_FORMATS: Dict[str, BedrockGuardrailImageFormat] = { + "png": "png", + "jpeg": "jpeg", + "jpg": "jpeg", +} + class QualifiedTextBlock(NamedTuple): """A piece of message text paired with its Bedrock grounding qualifier (if any).""" @@ -273,28 +284,101 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): ) return cleaned or None - def _create_bedrock_input_content_request(self, messages: Optional[List[AllMessageValues]]) -> BedrockRequest: + async def _create_bedrock_input_content_request(self, messages: Optional[List[AllMessageValues]]) -> BedrockRequest: """ Create a bedrock request for the input content - the LLM request. + + Text and image parts are both sent, so a guardrail with the IMAGE modality + enabled inspects the image the caller actually sent instead of only the text + that happened to sit next to it. """ bedrock_request: BedrockRequest = BedrockRequest(source="INPUT") - bedrock_request_content: List[BedrockContentItem] = [] if messages is None: return bedrock_request - for message in messages: - blocks = self.get_content_items_for_message(message=message) - if blocks is None: - continue - for block in blocks: - # INPUT scans send plain text only. Grounding qualifiers are attached - # exclusively when assembling the OUTPUT request, so a caller cannot use - # a grounding_source/query tag to change how input-safety policies treat - # their content (which would be an input-guardrail bypass). - bedrock_request_content.append(BedrockContentItem(text=BedrockTextContent(text=block.text))) - bedrock_request["content"] = bedrock_request_content + per_message = await asyncio.gather(*(self._build_input_content_items(message=message) for message in messages)) + bedrock_request["content"] = [item for items in per_message for item in items] return bedrock_request + async def _build_input_content_items(self, message: AllMessageValues) -> Tuple[BedrockContentItem, ...]: + """Flatten one request message into ApplyGuardrail INPUT content items. + + INPUT scans send plain text only. Grounding qualifiers are attached + exclusively when assembling the OUTPUT request, so a caller cannot use a + grounding_source/query tag to change how input-safety policies treat their + content (which would be an input-guardrail bypass). + """ + content = message.get("content") + if content is None: + return () + if isinstance(content, str): + return (BedrockContentItem(text=BedrockTextContent(text=content)),) + if not isinstance(content, list): + return () + parts = cast(Tuple[object, ...], tuple(content)) + items = await asyncio.gather(*(self._build_input_content_item(item=item) for item in parts)) + return tuple(item for item in items if item is not None) + + async def _build_input_content_item(self, item: object) -> Optional[BedrockContentItem]: + if isinstance(item, str): + return BedrockContentItem(text=BedrockTextContent(text=item)) + if not isinstance(item, dict): + return None + part = cast(Mapping[str, object], item) + text = part.get("text") + if isinstance(text, str): + return BedrockContentItem(text=BedrockTextContent(text=text)) + image_url = self._get_image_url(item=part) + if image_url is None: + return None + return await self._build_image_content_item(image_url=image_url) + + @staticmethod + def _get_image_url(item: Mapping[str, object]) -> Optional[str]: + if item.get("type") != "image_url": + return None + image_url = item.get("image_url") + if isinstance(image_url, str): + return image_url + if isinstance(image_url, dict): + url = cast(Mapping[str, object], image_url).get("url") + return url if isinstance(url, str) else None + return None + + async def _build_image_content_item(self, image_url: str) -> Optional[BedrockContentItem]: + """Decode/fetch an image part into an ApplyGuardrail image block. + + Anything ApplyGuardrail cannot take (unsupported format, unreachable url) is + skipped with a warning instead of failing the call, so an unscannable + attachment does not take the whole request down. + """ + try: + block = await BedrockImageProcessor.process_image_async(image_url=image_url, format=None) + except Exception as e: + verbose_proxy_logger.warning( + "Bedrock Guardrail %s: could not read image content, it will not be scanned: %s", + self.guardrail_name, + str(e), + ) + return None + + image_block = block.get("image") + image_format = _APPLY_GUARDRAIL_IMAGE_FORMATS.get(str(image_block.get("format"))) if image_block else None + image_bytes = image_block.get("source", {}).get("bytes") if image_block else None + if image_format is None or not image_bytes: + verbose_proxy_logger.warning( + "Bedrock Guardrail %s: attachment is not a png/jpeg image, ApplyGuardrail cannot scan it; skipping", + self.guardrail_name, + ) + return None + + return BedrockContentItem( + image=BedrockImageContent( + format=image_format, + source=BedrockGuardrailImageSource(bytes=image_bytes), + ) + ) + def _create_bedrock_output_content_request( self, response: Union[Any, ModelResponse], @@ -343,7 +427,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): items.append(self._build_content_item(block)) return items - def convert_to_bedrock_format( + async def convert_to_bedrock_format( self, source: Literal["INPUT", "OUTPUT"], messages: Optional[List[AllMessageValues]] = None, @@ -360,7 +444,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): """ bedrock_request: BedrockRequest = BedrockRequest(source=source) if source == "INPUT": - bedrock_request = self._create_bedrock_input_content_request(messages=messages) + bedrock_request = await self._create_bedrock_input_content_request(messages=messages) elif source == "OUTPUT": bedrock_request = self._create_bedrock_output_content_request(response=response, messages=messages) return bedrock_request @@ -763,7 +847,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): start_time = datetime.now(timezone.utc) credentials, aws_region_name = self._load_credentials() bedrock_request_data: dict = dict( - self.convert_to_bedrock_format(source=source, messages=messages, response=response) + await self.convert_to_bedrock_format(source=source, messages=messages, response=response) ) bedrock_guardrail_response: BedrockGuardrailResponse = BedrockGuardrailResponse() api_key: Optional[str] = None @@ -1983,6 +2067,8 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): masking_index += 1 if item is not None: new_content.append(item) + else: + new_content.append(item) return new_content, masking_index diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/types/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index ad9c9b44857..0043b516460 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -12,8 +12,21 @@ class BedrockTextContent(TypedDict, total=False): qualifiers: List[BedrockGuardrailQualifier] +BedrockGuardrailImageFormat = Literal["png", "jpeg"] + + +class BedrockGuardrailImageSource(TypedDict, total=False): + bytes: str + + +class BedrockImageContent(TypedDict, total=False): + format: BedrockGuardrailImageFormat + source: BedrockGuardrailImageSource + + class BedrockContentItem(TypedDict, total=False): text: BedrockTextContent + image: BedrockImageContent class BedrockRequest(TypedDict, total=False): diff --git a/tests/guardrails_tests/test_bedrock_guardrails.py b/tests/guardrails_tests/test_bedrock_guardrails.py index 823ee05839f..d9983913d0a 100644 --- a/tests/guardrails_tests/test_bedrock_guardrails.py +++ b/tests/guardrails_tests/test_bedrock_guardrails.py @@ -944,7 +944,7 @@ async def test_convert_to_bedrock_format_input_source(): ] # Call the method - result = guardrail.convert_to_bedrock_format(source="INPUT", messages=mock_messages) + result = await guardrail.convert_to_bedrock_format(source="INPUT", messages=mock_messages) # Verify the result structure assert isinstance(result, dict) @@ -1006,7 +1006,7 @@ async def test_convert_to_bedrock_format_output_source(): ) # Call the method - result = guardrail.convert_to_bedrock_format( + result = await guardrail.convert_to_bedrock_format( source="OUTPUT", response=mock_response ) diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py index 53f32fd96fb..d3056b7b417 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py @@ -2,6 +2,7 @@ Unit tests for Bedrock Guardrails """ +import asyncio import json import os import sys @@ -2567,15 +2568,19 @@ _GUARD_BLOCK = { def _input_request(messages: list) -> dict: """Arrange a guardrail and act: build the Bedrock INPUT payload.""" - return _grounding_guardrail().convert_to_bedrock_format( - source="INPUT", messages=messages + return asyncio.run( + _grounding_guardrail().convert_to_bedrock_format( + source="INPUT", messages=messages + ) ) def _output_request(messages: list, response=None) -> dict: """Arrange a guardrail and act: build the Bedrock OUTPUT payload.""" - return _grounding_guardrail().convert_to_bedrock_format( - source="OUTPUT", response=response, messages=messages + return asyncio.run( + _grounding_guardrail().convert_to_bedrock_format( + source="OUTPUT", response=response, messages=messages + ) ) @@ -3614,3 +3619,107 @@ class TestBedrockIncrementalFlagInteractions: ) scanned = [m["content"] for m in mock_api.call_args.kwargs["messages"]] assert scanned == ["q1"], "latest-role selection must exclude the system prompt" + + +class TestBedrockGuardrailImageInput: + """Image parts must reach ApplyGuardrail, not just the text sitting next to them.""" + + _PNG_DATA_URI = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUg==" + + def _guardrail(self) -> BedrockGuardrail: + return BedrockGuardrail( + guardrail_name="bedrock-image", + guardrailIdentifier="gr-image", + guardrailVersion="DRAFT", + ) + + @pytest.mark.asyncio + async def test_inline_image_is_sent_for_scanning(self): + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "what does this say?"}, + {"type": "image_url", "image_url": {"url": self._PNG_DATA_URI}}, + ], + } + ] + + request = await self._guardrail().convert_to_bedrock_format( + source="INPUT", messages=messages + ) + + assert request["content"] == [ + {"text": {"text": "what does this say?"}}, + { + "image": { + "format": "png", + "source": {"bytes": self._PNG_DATA_URI.split(",")[1]}, + } + }, + ] + + @pytest.mark.asyncio + async def test_unscannable_image_is_skipped_without_failing_the_scan(self): + """ApplyGuardrail takes png/jpeg only; a gif is dropped instead of 400ing the call.""" + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "hello"}, + { + "type": "image_url", + "image_url": {"url": "data:image/gif;base64,R0lGODlhAQABAAAAACw="}, + }, + {"type": "image_url", "image_url": {"url": "not-an-image"}}, + ], + } + ] + + request = await self._guardrail().convert_to_bedrock_format( + source="INPUT", messages=messages + ) + + assert request["content"] == [{"text": {"text": "hello"}}] + + @pytest.mark.asyncio + async def test_remote_image_is_fetched_and_scanned(self): + messages = [ + { + "role": "user", + "content": [ + {"type": "image_url", "image_url": {"url": "https://example.com/a.jpg"}} + ], + } + ] + + with patch( + "litellm.litellm_core_utils.prompt_templates.factory.BedrockImageProcessor.get_image_details_async", + new=AsyncMock(return_value=("Zm9v", "image/jpeg")), + ): + request = await self._guardrail().convert_to_bedrock_format( + source="INPUT", messages=messages + ) + + assert request["content"] == [ + {"image": {"format": "jpeg", "source": {"bytes": "Zm9v"}}} + ] + + def test_masking_keeps_image_parts_in_the_request(self): + """Masking rewrites text in place; the image must survive to reach the model.""" + image_part = {"type": "image_url", "image_url": {"url": self._PNG_DATA_URI}} + messages = [ + { + "role": "user", + "content": [{"type": "text", "text": "my ssn is 123-45-6789"}, image_part], + } + ] + + updated = self._guardrail()._apply_masking_to_messages( + messages=messages, masked_texts=["my ssn is {SSN}"] + ) + + assert updated[0]["content"] == [ + {"type": "text", "text": "my ssn is {SSN}"}, + image_part, + ] diff --git a/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py b/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py index 1c452e2fb6c..61b868ae65d 100644 --- a/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py +++ b/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py @@ -798,7 +798,9 @@ async def test_bedrock_guardrail_make_api_request_passes_api_key(): guardrail_hook.async_handler, "post", AsyncMock(return_value=mock_response) ), patch.object(guardrail_hook, "_load_credentials") as mock_load_creds, - patch.object(guardrail_hook, "convert_to_bedrock_format") as mock_convert, + patch.object( + guardrail_hook, "convert_to_bedrock_format", new_callable=AsyncMock + ) as mock_convert, patch.object( guardrail_hook, "get_guardrail_dynamic_request_body_params" ) as mock_get_params,