From 410f54dc72389d59300a07e9c4fdb26a40628489 Mon Sep 17 00:00:00 2001 From: Chesars Date: Sat, 21 Mar 2026 23:06:24 -0300 Subject: [PATCH 1/5] 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 2371bc4865..a3d681ea51 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 0f6d85525d..7d6e3a1c8a 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 0000000000..e69de29bb2 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 0000000000..4a0df332f6 --- /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 From 806dd31158897018b7105b345fde48eb1b210572 Mon Sep 17 00:00:00 2001 From: Chesars Date: Sat, 21 Mar 2026 23:19:20 -0300 Subject: [PATCH 2/5] test: add multimodal mixed input test for process_response --- .../test_batch_embed_content_transformation.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) 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 index 4a0df332f6..2658646d59 100644 --- 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 @@ -177,3 +177,19 @@ class TestProcessResponse: assert result.data[0]["index"] == 0 assert result.data[1]["index"] == 1 assert result.data[2]["index"] == 2 + + def test_multimodal_mixed_input(self): + """process_response works with mixed text + multimodal inputs.""" + predictions: VertexAIBatchEmbeddingsResponseObject = { + "embeddings": [{"values": [0.1, 0.2]}, {"values": [0.3, 0.4]}] + } + result = process_response( + input=["hello", IMAGE_DATA_URI], + model_response=EmbeddingResponse(), + model="gemini-embedding-2-preview", + _predictions=predictions, + ) + assert len(result.data) == 2 + assert result.data[0]["index"] == 0 + assert result.data[1]["index"] == 1 + assert result.usage.prompt_tokens >= 0 From 9c4c34177eb1da358644ff73e8998cf6c7964a51 Mon Sep 17 00:00:00 2001 From: Chesars Date: Sat, 21 Mar 2026 23:23:41 -0300 Subject: [PATCH 3/5] refactor: reuse _build_part_for_input in embed_content transform --- .../batch_embed_content_transformation.py | 24 +------------------ 1 file changed, 1 insertion(+), 23 deletions(-) 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 7d6e3a1c8a..39de34169c 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 @@ -238,29 +238,7 @@ def transform_openai_input_gemini_embed_content( for element in input_list: if not isinstance(element, str): raise ValueError(f"Unsupported input type: {type(element)}") - - 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} - parts.append(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, - } - parts.append(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"], - } - parts.append(PartType(file_data=file_data_ref)) - else: - parts.append(PartType(text=element)) + parts.append(_build_part_for_input(element, resolved_files=resolved_files)) request_body: dict = { "content": ContentType(parts=parts), From bb247685da9d9c95e6a3a5b61047e3222f2f91cd Mon Sep 17 00:00:00 2001 From: Chesars Date: Sat, 21 Mar 2026 23:29:45 -0300 Subject: [PATCH 4/5] fix: skip token counting for multimodal inputs in process_response --- .../batch_embed_content_transformation.py | 7 +++++-- .../test_batch_embed_content_transformation.py | 2 +- 2 files changed, 6 insertions(+), 3 deletions(-) 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 39de34169c..834d3a5c0a 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 @@ -312,8 +312,11 @@ def process_response( model_response.data = openai_embeddings model_response.model = model - input_text = get_formatted_prompt(data={"input": input}, call_type="embedding") - prompt_tokens = token_counter(model=model, text=input_text) + if _is_multimodal_input(input): + prompt_tokens = 0 + else: + input_text = get_formatted_prompt(data={"input": input}, call_type="embedding") + prompt_tokens = token_counter(model=model, text=input_text) model_response.usage = Usage( prompt_tokens=prompt_tokens, total_tokens=prompt_tokens ) 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 index 2658646d59..1417c04503 100644 --- 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 @@ -192,4 +192,4 @@ class TestProcessResponse: assert len(result.data) == 2 assert result.data[0]["index"] == 0 assert result.data[1]["index"] == 1 - assert result.usage.prompt_tokens >= 0 + assert result.usage.prompt_tokens == 0 From 883e150804342ece5152feebade0f37b811ae5a3 Mon Sep 17 00:00:00 2001 From: Chesars Date: Sat, 21 Mar 2026 23:37:34 -0300 Subject: [PATCH 5/5] fix: count text tokens only for mixed multimodal inputs, add embed_content tests --- .../batch_embed_content_transformation.py | 14 ++++++- ...test_batch_embed_content_transformation.py | 38 ++++++++++++++++++- 2 files changed, 50 insertions(+), 2 deletions(-) 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 834d3a5c0a..ca1c5d1f58 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 @@ -313,7 +313,19 @@ def process_response( model_response.model = model if _is_multimodal_input(input): - prompt_tokens = 0 + input_list = input if isinstance(input, list) else [input] + text_elements = [ + e for e in input_list + if isinstance(e, str) + and not (e.startswith("data:") and ";base64," in e) + and not _is_gcs_url(e) + and not _is_file_reference(e) + ] + if text_elements: + input_text = get_formatted_prompt(data={"input": text_elements}, call_type="embedding") + prompt_tokens = token_counter(model=model, text=input_text) + else: + prompt_tokens = 0 else: input_text = get_formatted_prompt(data={"input": input}, call_type="embedding") prompt_tokens = token_counter(model=model, text=input_text) 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 index 1417c04503..c15da3cdee 100644 --- 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 @@ -20,6 +20,7 @@ from litellm.llms.vertex_ai.gemini_embeddings.batch_embed_content_transformation _is_multimodal_input, process_response, transform_openai_input_gemini_content, + transform_openai_input_gemini_embed_content, ) from litellm.types.llms.vertex_ai import VertexAIBatchEmbeddingsResponseObject from litellm.types.utils import EmbeddingResponse @@ -140,6 +141,40 @@ class TestTransformOpenaiInputGeminiContent: assert len(result["requests"]) == 3 +class TestTransformOpenaiInputGeminiEmbedContent: + """Test transform_openai_input_gemini_embed_content (vertex_ai / embedContent path).""" + + def test_text_and_image_combined(self): + result = transform_openai_input_gemini_embed_content( + input=["hello", IMAGE_DATA_URI], + model="gemini-embedding-2-preview", + optional_params={}, + ) + assert "content" in result + parts = result["content"]["parts"] + assert len(parts) == 2 + assert parts[0]["text"] == "hello" + assert parts[1]["inline_data"] is not None + + def test_gcs_url(self): + result = transform_openai_input_gemini_embed_content( + input=[GCS_URL], + model="gemini-embedding-2-preview", + optional_params={}, + ) + parts = result["content"]["parts"] + assert len(parts) == 1 + assert parts[0]["file_data"]["file_uri"] == GCS_URL + + def test_dimensions_mapped(self): + result = transform_openai_input_gemini_embed_content( + input="hello", + model="gemini-embedding-2-preview", + optional_params={"dimensions": 256}, + ) + assert result["outputDimensionality"] == 256 + + class TestProcessResponse: """Test that process_response sets correct indices.""" @@ -192,4 +227,5 @@ class TestProcessResponse: assert len(result.data) == 2 assert result.data[0]["index"] == 0 assert result.data[1]["index"] == 1 - assert result.usage.prompt_tokens == 0 + # Should count tokens only for the text element, not the image + assert result.usage.prompt_tokens > 0