mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-06 08:16:43 +00:00
fix(guardrails): close patch-coverage gaps in bedrock image handling
Codecov flagged 7 uncovered lines added by this PR. Two were real test gaps: the remote-url rejection branch in _build_image_content_item and the non-dict-part skip in _file_backed_parts. Both get direct tests now. The other five were dead code: _handle_unscannable_image always raises, so the `return None` that followed each call in _build_image_content_item could never execute. Typed it NoReturn, dropped the unreachable returns, and tightened _build_image_content_item's return type from `BedrockContentItem | None` to `BedrockContentItem` since every remaining path either raises or returns a real item.
This commit is contained in:
parent
a46896b436
commit
24359ce82d
2 changed files with 37 additions and 7 deletions
|
|
@ -20,7 +20,7 @@ from collections.abc import AsyncGenerator, Mapping, Sequence
|
|||
from datetime import datetime, timezone
|
||||
from itertools import accumulate, groupby
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, NamedTuple, Optional, cast
|
||||
from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, NamedTuple, NoReturn, Optional, cast
|
||||
|
||||
import httpx
|
||||
from fastapi import HTTPException
|
||||
|
|
@ -492,7 +492,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
)
|
||||
return count
|
||||
|
||||
def _handle_unscannable_image(self, reason: str) -> None:
|
||||
def _handle_unscannable_image(self, reason: str) -> NoReturn:
|
||||
"""Block an image part ApplyGuardrail cannot scan.
|
||||
|
||||
The image reaches the model either way, so skipping it silently would let a
|
||||
|
|
@ -594,7 +594,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
found += 1
|
||||
return found
|
||||
|
||||
async def _build_image_content_item(self, image_url: str) -> BedrockContentItem | None:
|
||||
async def _build_image_content_item(self, image_url: str) -> BedrockContentItem:
|
||||
"""Decode an inline image into an ApplyGuardrail image block.
|
||||
|
||||
A remote url is named as its own rejection rather than left to the decoder:
|
||||
|
|
@ -605,13 +605,11 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
"""
|
||||
if image_url.startswith(("http://", "https://")):
|
||||
self._handle_unscannable_image(reason="remote image URLs are not supported")
|
||||
return None
|
||||
|
||||
try:
|
||||
block: Final = await BedrockImageProcessor.process_image_async(image_url=image_url, format=None)
|
||||
except (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 = (
|
||||
|
|
@ -621,7 +619,6 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
image_bytes: Final = image_source.get("bytes") if image_source 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
|
||||
|
||||
# base64 decodes to roughly 3/4 of its length; estimate rather than decode the
|
||||
# whole image a second time just to measure it.
|
||||
|
|
@ -630,7 +627,6 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
self._handle_unscannable_image(
|
||||
reason=f"image is {decoded_size / 1024 / 1024:.1f} MB, over ApplyGuardrail's 4 MB limit"
|
||||
)
|
||||
return None
|
||||
|
||||
return BedrockContentItem(
|
||||
image=BedrockImageContent(
|
||||
|
|
|
|||
|
|
@ -5551,6 +5551,40 @@ class TestBedrockGuardrailImageInput:
|
|||
|
||||
assert "could not be read" in str(exc_info.value.detail) or "not a png/jpeg" in str(exc_info.value.detail)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_remote_image_url_is_rejected_without_being_fetched(self):
|
||||
"""A remote url is named as its own rejection rather than left to the decoder.
|
||||
|
||||
Fetching one safely (size cap, SSRF/redirect validation) is separate work; this
|
||||
PR only scans inline images, so a url has to fail closed rather than be ignored
|
||||
or silently forwarded to the model unscanned.
|
||||
"""
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [{"type": "image_url", "image_url": {"url": "https://example.com/pic.png"}}],
|
||||
}
|
||||
]
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await self._guardrail().convert_to_bedrock_format(source="INPUT", messages=messages)
|
||||
|
||||
assert "remote image URLs are not supported" in str(exc_info.value.detail)
|
||||
|
||||
def test_file_backed_part_counting_skips_non_mapping_entries(self):
|
||||
"""A content list may mix plain strings in with typed parts.
|
||||
|
||||
`_file_backed_image_count` walks the raw request rather than a normalized
|
||||
shape, so a non-dict entry must be skipped rather than raise or be miscounted.
|
||||
"""
|
||||
content = [
|
||||
"just a string",
|
||||
{"type": "text", "text": "hi"},
|
||||
{"type": "image", "source": {"type": "file", "file_id": "file_abc"}},
|
||||
]
|
||||
|
||||
assert BedrockGuardrail._file_backed_parts(content) == 1
|
||||
|
||||
def test_the_url_helper_guards_its_own_inputs(self):
|
||||
"""Exercised directly so the guards are not dropped in a later refactor."""
|
||||
assert BedrockGuardrail._get_image_url(item={"type": "image_url"}) is None
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue