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():