diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index c70a2ee8a74..62454b89532 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -10,6 +10,7 @@ import sys sys.path.insert(0, os.path.abspath("../..")) # Adds the parent directory to the system path import asyncio +import binascii import copy import json import re @@ -32,6 +33,8 @@ from litellm.exceptions import ModifyResponseException from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.litellm_core_utils.core_helpers import redact_nested_match_and_regex_keys from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import bedrock_guardrail_cost +from litellm.litellm_core_utils.prompt_templates.factory import BedrockImageProcessor +from litellm.litellm_core_utils.url_utils import SSRFError from litellm.llms.anthropic.chat.guardrail_translation.handler import AnthropicMessagesHandler from litellm.llms.base_llm.guardrail_translation.utils import ( effective_scan_only_tool_results_for_guardrail, @@ -59,10 +62,13 @@ from litellm.types.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( BedrockChecksViolation, BedrockContentItem, BedrockGuardrailChecksResponse, + BedrockGuardrailImageFormat, + BedrockGuardrailImageSource, BedrockGuardrailOutput, BedrockGuardrailQualifier, BedrockGuardrailResponse, BedrockGuardrailUsage, + BedrockImageContent, BedrockRequest, BedrockTextContent, ) @@ -133,6 +139,16 @@ _CONTENT_TYPE_TO_QUALIFIER: Final[dict[str, BedrockGuardrailQualifier]] = { # must not be graded against as if it were the application's own source material. _GROUNDING_SOURCE_TRUSTED_ROLES: Final = frozenset({"system", "developer"}) +# ApplyGuardrail only accepts png/jpeg image blocks. Anything else (gif, webp, ...) +# has no representation in the payload, so it cannot be scanned at all. +_APPLY_GUARDRAIL_IMAGE_FORMATS: Final[ + dict[str, BedrockGuardrailImageFormat] +] = { # mutable-ok: module-level lookup table, never mutated + "png": "png", + "jpeg": "jpeg", + "jpg": "jpeg", +} + class QualifiedTextBlock(NamedTuple): """A piece of message text paired with its Bedrock grounding qualifier (if any).""" @@ -221,6 +237,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): prompt_attack_threshold: float | None = 0.5, pii_confidence_threshold: float | None = 0.5, chunk_budget_chars: int = BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS, + on_unscannable_image: Literal["block", "allow"] = "block", **kwargs, ): self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback) @@ -229,6 +246,10 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): self.guardrail_provider = "bedrock" self.chunk_budget_chars = chunk_budget_chars self.experimental_use_latest_role_message_only = bool(kwargs.get("experimental_use_latest_role_message_only")) + # What to do with an image part ApplyGuardrail cannot scan (non png/jpeg, or a + # remote url we refuse to fetch). Defaults to blocking: the image reaches the + # model regardless, so allowing it would be a silent guardrail bypass. + self.on_unscannable_image: Literal["block", "allow"] = on_unscannable_image # Resource-less, detect-only InvokeGuardrailChecks mode. Present `checks` # routes the guardrail to InvokeGuardrailChecks; absent => ApplyGuardrail. @@ -316,28 +337,139 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): ) return cleaned or None - def _create_bedrock_input_content_request(self, messages: list[AllMessageValues] | None) -> BedrockRequest: + async def _create_bedrock_input_content_request(self, messages: list[AllMessageValues] | None) -> 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: Final[BedrockRequest] = BedrockRequest(source="INPUT") - bedrock_request_content: Final[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)) + # mutable-ok: BedrockRequest["content"] is a list in the AWS wire format + 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 text and image parts. 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: Final = cast( # cast-ok: AllMessageValues content is a union of part TypedDicts + tuple[object, ...], tuple(content) + ) + items: Final = 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) -> BedrockContentItem | None: + if isinstance(item, str): + return BedrockContentItem(text=BedrockTextContent(text=item)) + if not isinstance(item, dict): + return None + part: Final = cast(Mapping[str, object], item) # cast-ok: narrowed to dict on the line above + text: Final = part.get("text") + if isinstance(text, str): + return BedrockContentItem(text=BedrockTextContent(text=text)) + image_url: Final = 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]) -> str | None: + if item.get("type") != "image_url": + return None + image_url: Final = item.get("image_url") + if isinstance(image_url, str): + return image_url + if isinstance(image_url, dict): + url: Final = cast(Mapping[str, object], image_url).get("url") # cast-ok: narrowed to dict on the line above + return url if isinstance(url, str) else None + return None + + def _handle_unscannable_image(self, reason: str) -> None: + """Decide what to do with an image part ApplyGuardrail cannot scan. + + The image still reaches the model either way, so skipping it silently would + let a caller defeat an IMAGE-modality guardrail just by picking a format the + API does not accept. `on_unscannable_image` defaults to "block" for that + reason; "allow" restores the permissive behavior for deployments that would + rather serve the request than fail it. + + Returns None so callers can `return self._handle_unscannable_image(...)`. + """ + if self.on_unscannable_image == "block": + raise HTTPException( + status_code=400, + detail={ # mutable-ok: HTTPException detail payload, serialized immediately + "error": "Violated guardrail policy", + "bedrock_guardrail_response": ( + f"Request contains an image the guardrail cannot scan ({reason}). " + "ApplyGuardrail accepts png/jpeg images only. Set " + "'on_unscannable_image: allow' on this guardrail to send such " + "requests to the model unscanned." + ), + "guardrail_name": self.guardrail_name, + }, + ) + verbose_proxy_logger.warning( + "Bedrock Guardrail %s: image part will not be scanned (%s); on_unscannable_image=allow, forwarding it to the model anyway", + self.guardrail_name, + reason, + ) + + async def _build_image_content_item(self, image_url: str) -> BedrockContentItem | None: + """Decode/fetch an image part into an ApplyGuardrail image block. + + Remote URLs are only fetched while LiteLLM's URL validation is on. With + validation disabled `async_safe_get` degrades to an unrestricted, redirect + following GET, and the URL comes straight from the caller, so fetching here + would turn the guardrail into an SSRF primitive. Such an image is treated as + unscannable instead. + """ + if not image_url.startswith("data:") and not getattr(litellm, "user_url_validation", True): + self._handle_unscannable_image( + reason=f"remote image url not fetched because litellm.user_url_validation is disabled: {image_url}" + ) + return None + + try: + block: Final = await BedrockImageProcessor.process_image_async(image_url=image_url, format=None) + except (httpx.HTTPError, SSRFError, ValueError, TypeError, KeyError, binascii.Error) as e: + self._handle_unscannable_image(reason=f"image content could not be read: {e}") + return None + + image_block: Final = block.get("image") + image_format: Final = ( + _APPLY_GUARDRAIL_IMAGE_FORMATS.get(str(image_block.get("format"))) if image_block else None + ) + # mutable-ok: {} is an empty default for .get, never mutated + image_bytes: Final = image_block.get("source", {}).get("bytes") if image_block else None + if image_format is None or not image_bytes: + self._handle_unscannable_image(reason="attachment is not a png/jpeg image") + return None + + return BedrockContentItem( + image=BedrockImageContent( + format=image_format, + source=BedrockGuardrailImageSource(bytes=image_bytes), + ) + ) + def _create_bedrock_output_content_request( self, response: Any | ModelResponse, @@ -386,7 +518,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: list[AllMessageValues] | None = None, @@ -403,7 +535,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 @@ -839,7 +971,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): """ start_time: Final = datetime.now(timezone.utc) bedrock_request_data: Final[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) ) api_key: str | None = None if request_data: @@ -2880,6 +3012,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/guardrails.py b/litellm/types/guardrails.py index c7cdfaad780..0223d651d4a 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -535,6 +535,13 @@ class BedrockGuardrailConfigModel(BaseModel): "still rejects is bisected automatically, so this value only trades round trips against " "batch size and cannot fail a request on its own.", ) + on_unscannable_image: Literal["block", "allow"] = Field( + default="block", + description="What to do with an image the guardrail cannot scan - ApplyGuardrail " + "accepts png/jpeg only, and remote image URLs are not fetched while " + "litellm.user_url_validation is disabled. 'block' (default) rejects the request; " + "'allow' logs a warning and sends the image to the model unscanned.", + ) class LakeraV2GuardrailConfigModel(BaseModel): diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/types/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index 8d66b624341..e194f4f1132 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -1,6 +1,6 @@ from typing import Literal -from typing_extensions import TypedDict +from typing_extensions import ReadOnly, TypedDict # Bedrock contextual grounding tags each content block so the guardrail knows # which text is the reference source, the user question, and the content to grade. @@ -12,8 +12,21 @@ class BedrockTextContent(TypedDict, total=False): qualifiers: list[BedrockGuardrailQualifier] +BedrockGuardrailImageFormat = Literal["png", "jpeg"] + + +class BedrockGuardrailImageSource(TypedDict, total=False): + bytes: ReadOnly[str] + + +class BedrockImageContent(TypedDict, total=False): + format: ReadOnly[BedrockGuardrailImageFormat] + source: ReadOnly[BedrockGuardrailImageSource] + + class BedrockContentItem(TypedDict, total=False): text: BedrockTextContent + image: ReadOnly[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 43d088268eb..bb958bcc304 100644 --- a/tests/guardrails_tests/test_bedrock_guardrails.py +++ b/tests/guardrails_tests/test_bedrock_guardrails.py @@ -945,7 +945,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) @@ -1007,7 +1007,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 dd339d4e51f..5eb7ff867ee 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,8 @@ Unit tests for Bedrock Guardrails """ +import asyncio +import base64 import json import sys from unittest.mock import AsyncMock, MagicMock, patch @@ -2408,12 +2410,14 @@ _GUARD_BLOCK = {"text": {"text": _GROUNDING_RESPONSE_TEXT, "qualifiers": ["guard 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) + ) def test_grounding_input_strips_grounding_and_query_qualifiers(): @@ -2961,9 +2965,7 @@ async def test_streaming_hook_reraises_guardrail_service_failures(): guardrail = _sse_guardrail() with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: - mock_api.side_effect = HTTPException( - status_code=500, detail="Bedrock guardrail throttle retries exhausted" - ) + mock_api.side_effect = HTTPException(status_code=500, detail="Bedrock guardrail throttle retries exhausted") with pytest.raises(HTTPException) as exc: await _drain_streaming_hook(guardrail) @@ -5274,3 +5276,169 @@ async def test_terminal_failure_logs_usage_and_cost_of_prior_passed_chunks(monke assert logged["guardrail_cost"] == pytest.approx(0.0003) assert logged["guardrail_response"]["usage"] == {"contentPolicyUnits": 2, "wordPolicyUnits": 1} assert "error" in logged["guardrail_response"] + + +class TestBedrockGuardrailImageInput: + """Image parts must reach ApplyGuardrail, not just the text sitting next to them.""" + + _PNG_DATA_URI = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUg==" + _GIF_DATA_URI = "data:image/gif;base64,R0lGODlhAQABAAAAACw=" + _JPEG_BYTES = b"\xff\xd8\xff\xdb" + # A literal, globally-routable IP keeps validate_url's getaddrinfo off DNS; the + # transport is faked in every test below, so no request leaves the process. + _REMOTE_IMAGE_URL = "https://93.184.216.34/a.jpg" + + def _jpeg_response(self, url: str) -> httpx.Response: + return httpx.Response( + 200, + content=self._JPEG_BYTES, + headers={"content-type": "image/jpeg"}, + request=httpx.Request("GET", url), + ) + + def _guardrail(self, **kwargs) -> BedrockGuardrail: + return BedrockGuardrail( + guardrail_name="bedrock-image", + guardrailIdentifier="gr-image", + guardrailVersion="DRAFT", + **kwargs, + ) + + @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_blocks_the_request_by_default(self): + """ApplyGuardrail takes png/jpeg only, and the image reaches the model either way. + + Skipping it silently would let a caller defeat an IMAGE-modality guardrail by + sending a gif, so the default is to reject the request. + """ + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "hello"}, + {"type": "image_url", "image_url": {"url": self._GIF_DATA_URI}}, + ], + } + ] + + with pytest.raises(HTTPException) as exc_info: + await self._guardrail().convert_to_bedrock_format(source="INPUT", messages=messages) + + assert exc_info.value.status_code == 400 + assert exc_info.value.detail["error"] == "Violated guardrail policy" + + @pytest.mark.asyncio + async def test_undecodable_image_blocks_the_request_by_default(self): + messages = [ + { + "role": "user", + "content": [{"type": "image_url", "image_url": {"url": "not-an-image"}}], + } + ] + + with pytest.raises(HTTPException) as exc_info: + await self._guardrail().convert_to_bedrock_format(source="INPUT", messages=messages) + + assert exc_info.value.status_code == 400 + + @pytest.mark.asyncio + async def test_unscannable_image_is_skipped_when_explicitly_allowed(self): + """on_unscannable_image: allow restores the permissive behavior, opt-in only.""" + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "hello"}, + {"type": "image_url", "image_url": {"url": self._GIF_DATA_URI}}, + {"type": "image_url", "image_url": {"url": "not-an-image"}}, + ], + } + ] + + request = await self._guardrail(on_unscannable_image="allow").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": self._REMOTE_IMAGE_URL}}], + } + ] + get = AsyncMock(return_value=self._jpeg_response(self._REMOTE_IMAGE_URL)) + + with patch.object(httpx.AsyncClient, "get", new=get): + request = await self._guardrail().convert_to_bedrock_format(source="INPUT", messages=messages) + + assert request["content"] == [ + { + "image": { + "format": "jpeg", + "source": {"bytes": base64.b64encode(self._JPEG_BYTES).decode()}, + } + } + ] + + @pytest.mark.asyncio + async def test_remote_image_is_not_fetched_when_url_validation_is_disabled(self, monkeypatch): + """With validation off, async_safe_get is an unrestricted redirect-following GET. + + The url comes straight from the caller, so fetching it here would make the + guardrail an SSRF primitive. Treat the image as unscannable instead, and make + no request at all. + """ + monkeypatch.setattr(litellm, "user_url_validation", False, raising=False) + url = "http://169.254.169.254/latest/meta-data/" + messages = [{"role": "user", "content": [{"type": "image_url", "image_url": {"url": url}}]}] + get = AsyncMock(return_value=self._jpeg_response(url)) + + with patch.object(httpx.AsyncClient, "get", new=get): + with pytest.raises(HTTPException): + await self._guardrail().convert_to_bedrock_format(source="INPUT", messages=messages) + + request = await self._guardrail(on_unscannable_image="allow").convert_to_bedrock_format( + source="INPUT", messages=messages + ) + + get.assert_not_awaited() + assert request["content"] == [] + + 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 45f5afef1bc..d6f7d5305be 100644 --- a/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py +++ b/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py @@ -896,7 +896,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,