From 410f54dc72389d59300a07e9c4fdb26a40628489 Mon Sep 17 00:00:00 2001 From: Chesars Date: Sat, 21 Mar 2026 23:06:24 -0300 Subject: [PATCH] fix(gemini): return separate embeddings for multimodal batch inputs (#24209) When multiple inputs were passed to the Gemini embedding endpoint and any contained multimodal data (images, audio, etc.), LiteLLM incorrectly used the `embedContent` endpoint which combines all inputs into a single aggregated embedding. Now uses `batchEmbedContents` with each input as a separate request, returning N embeddings for N inputs as expected. Also fixes hardcoded index=0 in batch embedding responses. --- .../batch_embed_content_handler.py | 23 ++- .../batch_embed_content_transformation.py | 63 ++++-- .../vertex_ai/gemini_embeddings/__init__.py | 0 ...test_batch_embed_content_transformation.py | 179 ++++++++++++++++++ 4 files changed, 247 insertions(+), 18 deletions(-) create mode 100644 tests/test_litellm/llms/vertex_ai/gemini_embeddings/__init__.py create mode 100644 tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py diff --git a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py index 2371bc4865a..a3d681ea515 100644 --- a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py +++ b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py @@ -151,8 +151,7 @@ class GoogleBatchEmbeddings(VertexLLM): optional_params = optional_params or {} - is_multimodal = _is_multimodal_input(input) - use_embed_content = is_multimodal or (custom_llm_provider == "vertex_ai") + use_embed_content = custom_llm_provider == "vertex_ai" mode: Literal["embedding", "batch_embedding"] if use_embed_content: mode = "embedding" @@ -215,8 +214,16 @@ class GoogleBatchEmbeddings(VertexLLM): resolved_files=resolved_files, ) else: + resolved_files = {} + if api_key and _is_multimodal_input(input): + resolved_files = self._resolve_file_references( + input=input, api_key=api_key, sync_handler=sync_handler + ) request_data = transform_openai_input_gemini_content( - input=input, model=model, optional_params=optional_params + input=input, + model=model, + optional_params=optional_params, + resolved_files=resolved_files, ) ## LOGGING @@ -303,8 +310,16 @@ class GoogleBatchEmbeddings(VertexLLM): resolved_files=resolved_files, ) else: + resolved_files = {} + if api_key and _is_multimodal_input(input): + resolved_files = await self._async_resolve_file_references( + input=input, api_key=api_key, async_handler=async_handler + ) data = transform_openai_input_gemini_content( - input=input, model=model, optional_params=optional_params or {} + input=input, + model=model, + optional_params=optional_params or {}, + resolved_files=resolved_files, ) ## LOGGING diff --git a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py index 0f6d85525d9..7d6e3a1c8ab 100644 --- a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py +++ b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py @@ -141,11 +141,51 @@ def _is_multimodal_input(input: EmbeddingInput) -> bool: return False +def _build_part_for_input( + element: str, + resolved_files: Optional[Dict[str, Dict[str, str]]] = None, +) -> PartType: + """ + Build a single PartType for an input element, handling text, data URIs, + file references, and GCS URLs. + """ + resolved_files = resolved_files or {} + + if element.startswith("data:") and ";base64," in element: + mime_type, base64_data = _parse_data_url(element) + blob: BlobType = {"mime_type": mime_type, "data": base64_data} + return PartType(inline_data=blob) + elif _is_gcs_url(element): + mime_type = _infer_mime_type_from_gcs_url(element) + file_data: FileDataType = { + "mime_type": mime_type, + "file_uri": element, + } + return PartType(file_data=file_data) + elif _is_file_reference(element): + if element not in resolved_files: + raise ValueError(f"File reference {element} not resolved") + file_info = resolved_files[element] + file_data_ref: FileDataType = { + "mime_type": file_info["mime_type"], + "file_uri": file_info["uri"], + } + return PartType(file_data=file_data_ref) + else: + return PartType(text=element) + + def transform_openai_input_gemini_content( - input: EmbeddingInput, model: str, optional_params: dict + input: EmbeddingInput, + model: str, + optional_params: dict, + resolved_files: Optional[Dict[str, Dict[str, str]]] = None, ) -> VertexAIBatchEmbeddingsRequestBody: """ - The content to embed. Only the parts.text fields will be counted. + Transform OpenAI embedding input to Gemini batchEmbedContents format. + + Each input element becomes a separate EmbedContentRequest, supporting + text, data URIs, file references, and GCS URLs. """ gemini_model_name = "models/{}".format(model) @@ -153,22 +193,17 @@ def transform_openai_input_gemini_content( if "dimensions" in gemini_params: gemini_params["outputDimensionality"] = gemini_params.pop("dimensions") + input_list = [input] if isinstance(input, str) else input requests: List[EmbedContentRequest] = [] - if isinstance(input, str): + + for element in input_list: + part = _build_part_for_input(element, resolved_files=resolved_files) request = EmbedContentRequest( model=gemini_model_name, - content=ContentType(parts=[PartType(text=input)]), + content=ContentType(parts=[part]), **gemini_params, ) requests.append(request) - else: - for i in input: - request = EmbedContentRequest( - model=gemini_model_name, - content=ContentType(parts=[PartType(text=i)]), - **gemini_params, - ) - requests.append(request) return VertexAIBatchEmbeddingsRequestBody(requests=requests) @@ -288,10 +323,10 @@ def process_response( _predictions: VertexAIBatchEmbeddingsResponseObject, ) -> EmbeddingResponse: openai_embeddings: List[Embedding] = [] - for embedding in _predictions["embeddings"]: + for idx, embedding in enumerate(_predictions["embeddings"]): openai_embedding = Embedding( embedding=embedding["values"], - index=0, + index=idx, object="embedding", ) openai_embeddings.append(openai_embedding) diff --git a/tests/test_litellm/llms/vertex_ai/gemini_embeddings/__init__.py b/tests/test_litellm/llms/vertex_ai/gemini_embeddings/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py b/tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py new file mode 100644 index 00000000000..4a0df332f63 --- /dev/null +++ b/tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py @@ -0,0 +1,179 @@ +""" +Tests for Gemini batchEmbedContents transformation logic. + +Covers: +- Text-only inputs (single and batch) +- Multimodal inputs (data URIs, GCS URLs, file references) +- Mixed text + multimodal inputs +- Response processing with correct indices +""" + +import os +import sys + +import pytest + +sys.path.insert(0, os.path.abspath("../../../../..")) + +from litellm.llms.vertex_ai.gemini_embeddings.batch_embed_content_transformation import ( + _build_part_for_input, + _is_multimodal_input, + process_response, + transform_openai_input_gemini_content, +) +from litellm.types.llms.vertex_ai import VertexAIBatchEmbeddingsResponseObject +from litellm.types.utils import EmbeddingResponse + + +IMAGE_DATA_URI = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAgAAAAIAQMAAAD+wSzIAAAABlBMVEX///+/v7+jQ3Y5AAAADklEQVQI12P4AIX8EAgALgAD/aNpbtEAAAAASUVORK5CYII" +GCS_URL = "gs://my-bucket/image.png" + + +class TestIsMultimodalInput: + def test_text_only_string(self): + assert _is_multimodal_input("hello world") is False + + def test_text_only_list(self): + assert _is_multimodal_input(["hello", "world"]) is False + + def test_data_uri(self): + assert _is_multimodal_input([IMAGE_DATA_URI]) is True + + def test_gcs_url(self): + assert _is_multimodal_input([GCS_URL]) is True + + def test_file_reference(self): + assert _is_multimodal_input(["files/abc123"]) is True + + def test_mixed_text_and_image(self): + assert _is_multimodal_input(["hello", IMAGE_DATA_URI]) is True + + +class TestBuildPartForInput: + def test_text_input(self): + part = _build_part_for_input("hello") + assert part["text"] == "hello" + assert part.get("inline_data") is None + + def test_data_uri_input(self): + part = _build_part_for_input(IMAGE_DATA_URI) + assert part.get("text") is None + assert part["inline_data"] is not None + assert part["inline_data"]["mime_type"] == "image/png" + + def test_gcs_url_input(self): + part = _build_part_for_input(GCS_URL) + assert part.get("text") is None + assert part["file_data"] is not None + assert part["file_data"]["mime_type"] == "image/png" + assert part["file_data"]["file_uri"] == GCS_URL + + def test_file_reference_resolved(self): + resolved = {"files/abc": {"mime_type": "image/jpeg", "uri": "https://example.com/abc"}} + part = _build_part_for_input("files/abc", resolved_files=resolved) + assert part["file_data"] is not None + assert part["file_data"]["mime_type"] == "image/jpeg" + + def test_file_reference_unresolved_raises(self): + with pytest.raises(ValueError, match="not resolved"): + _build_part_for_input("files/abc") + + +class TestTransformOpenaiInputGeminiContent: + """Test that transform_openai_input_gemini_content creates separate requests per input.""" + + def test_single_text(self): + result = transform_openai_input_gemini_content( + input="hello", model="gemini-embedding-2-preview", optional_params={} + ) + assert len(result["requests"]) == 1 + assert result["requests"][0]["content"]["parts"][0]["text"] == "hello" + + def test_multiple_texts(self): + result = transform_openai_input_gemini_content( + input=["hello", "world"], model="gemini-embedding-2-preview", optional_params={} + ) + assert len(result["requests"]) == 2 + assert result["requests"][0]["content"]["parts"][0]["text"] == "hello" + assert result["requests"][1]["content"]["parts"][0]["text"] == "world" + + def test_multimodal_inputs_are_separate_requests(self): + """Key regression test for #24209: each input becomes its own request.""" + result = transform_openai_input_gemini_content( + input=["The food was delicious", IMAGE_DATA_URI], + model="gemini-embedding-2-preview", + optional_params={}, + ) + assert len(result["requests"]) == 2 + # First request is text + assert result["requests"][0]["content"]["parts"][0]["text"] == "The food was delicious" + # Second request is image + assert result["requests"][1]["content"]["parts"][0]["inline_data"] is not None + + def test_dimensions_mapped_to_output_dimensionality(self): + result = transform_openai_input_gemini_content( + input="hello", + model="gemini-embedding-2-preview", + optional_params={"dimensions": 256}, + ) + assert result["requests"][0]["outputDimensionality"] == 256 + + def test_model_name_prefixed(self): + result = transform_openai_input_gemini_content( + input="hello", model="gemini-embedding-2-preview", optional_params={} + ) + assert result["requests"][0]["model"] == "models/gemini-embedding-2-preview" + + def test_gcs_url_input(self): + result = transform_openai_input_gemini_content( + input=[GCS_URL], model="gemini-embedding-2-preview", optional_params={} + ) + assert len(result["requests"]) == 1 + assert result["requests"][0]["content"]["parts"][0]["file_data"] is not None + + def test_mixed_text_image_gcs(self): + result = transform_openai_input_gemini_content( + input=["hello", IMAGE_DATA_URI, GCS_URL], + model="gemini-embedding-2-preview", + optional_params={}, + ) + assert len(result["requests"]) == 3 + + +class TestProcessResponse: + """Test that process_response sets correct indices.""" + + def test_single_embedding_index(self): + predictions: VertexAIBatchEmbeddingsResponseObject = { + "embeddings": [{"values": [0.1, 0.2]}] + } + model_response = EmbeddingResponse() + result = process_response( + input="hello", + model_response=model_response, + model="gemini-embedding-2-preview", + _predictions=predictions, + ) + assert len(result.data) == 1 + assert result.data[0]["index"] == 0 + + def test_multiple_embeddings_have_correct_indices(self): + """Regression test: indices should be 0, 1, 2... not all 0.""" + predictions: VertexAIBatchEmbeddingsResponseObject = { + "embeddings": [ + {"values": [0.1, 0.2]}, + {"values": [0.3, 0.4]}, + {"values": [0.5, 0.6]}, + ] + } + model_response = EmbeddingResponse() + result = process_response( + input=["a", "b", "c"], + model_response=model_response, + model="gemini-embedding-2-preview", + _predictions=predictions, + ) + assert len(result.data) == 3 + assert result.data[0]["index"] == 0 + assert result.data[1]["index"] == 1 + assert result.data[2]["index"] == 2