From 0a822e460cf71e6f14cdac511b21d40524001654 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 11 Nov 2026 00:39:14 -0800 Subject: [PATCH] fix(vector_stores): accept line-wrapped base64 in query_image data Bedrock takes inlineContent wrapped with \n or \r\n line breaks, as plain base64 writes it, so the check ignores only those before decoding and the data goes to Bedrock as sent (cherry picked from commit c2cad843b07609f3d216a46db6150c186568d595) --- litellm/vector_stores/search_query.py | 3 ++- tests/unit/vector_stores/test_main.py | 20 +++++++++++++++++--- 2 files changed, 19 insertions(+), 4 deletions(-) diff --git a/litellm/vector_stores/search_query.py b/litellm/vector_stores/search_query.py index 901f3889c7e..57770225889 100644 --- a/litellm/vector_stores/search_query.py +++ b/litellm/vector_stores/search_query.py @@ -15,6 +15,7 @@ from litellm.types.vector_stores import RETRIEVAL_QUERY_PARAMS, VectorStoreSearc _QUERY_IMAGE: Final = TypeAdapter(VectorStoreSearchQueryImage) _Query: Final = TypeVar("_Query", bound=str | Sequence[str]) _CALLER_QUERY_PARAMS: Final = ("query_image", *RETRIEVAL_QUERY_PARAMS) +_BASE64_LINE_BREAKS: Final = str.maketrans("", "", "\r\n") @dataclass(frozen=True, slots=True) @@ -80,7 +81,7 @@ def _problem(detail: ErrorDetails) -> str: def _is_base64(data: str) -> bool: try: - base64.b64decode(data, validate=True) + base64.b64decode(data.translate(_BASE64_LINE_BREAKS), validate=True) except ValueError: return False return True diff --git a/tests/unit/vector_stores/test_main.py b/tests/unit/vector_stores/test_main.py index f8d19707a85..58d29c808ed 100644 --- a/tests/unit/vector_stores/test_main.py +++ b/tests/unit/vector_stores/test_main.py @@ -209,16 +209,30 @@ def test_search_rejects_an_invalid_query_image_without_echoing_its_data(): assert image_data not in str(exc_info.value) -def test_search_rejects_query_image_data_that_is_not_base64(): +@pytest.mark.parametrize("data", ["not base64!", "/9j/4AAQ SkZJRg==", "/9j/4AAQ\tSkZJRg=="]) +def test_search_rejects_query_image_data_that_is_not_base64(data: str): with pytest.raises(litellm.BadRequestError) as exc_info: _search_with_recorded_post( custom_llm_provider="bedrock", - query_image={"type": "base64", "media_type": "image/jpeg", "data": "not base64!"}, + query_image={"type": "base64", "media_type": "image/jpeg", "data": data}, ) assert exc_info.value.status_code == 400 assert "query_image.data: not valid base64" in str(exc_info.value) - assert "not base64!" not in str(exc_info.value) + assert data not in str(exc_info.value) + + +@pytest.mark.parametrize("line_break", ["\n", "\r\n"]) +def test_search_sends_line_wrapped_query_image_data_to_bedrock_as_sent(line_break: str): + data: Final = line_break.join(("/9j/4AAQ", "SkZJRg==")) + + client: Final = _search_with_recorded_post( + custom_llm_provider="bedrock", + query_image={"type": "base64", "media_type": "image/jpeg", "data": data}, + ) + + posted: Final = json.loads(client.post.call_args.kwargs["data"]) + assert posted["retrievalQuery"]["image"]["inlineContent"] == data def test_search_by_a_query_image_sent_in_extra_body_sends_a_bedrock_image_retrieval_query():