Merge pull request #24337 from Chesars/fix/gemini-multimodal-batch-embeddings-24209
fix(gemini): return separate embeddings for multimodal inputs
This commit is contained in:
Cesar Garcia
GitHub
commit
a6462143be
@ -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
|
||||
|
||||
@ -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)
|
||||
|
||||
@ -155,22 +195,17 @@ def transform_openai_input_gemini_content(
|
||||
if "task_type" in gemini_params:
|
||||
gemini_params["taskType"] = gemini_params.pop("task_type")
|
||||
|
||||
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)
|
||||
|
||||
@ -207,29 +242,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),
|
||||
@ -292,10 +305,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)
|
||||
@ -303,8 +316,23 @@ 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):
|
||||
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)
|
||||
model_response.usage = Usage(
|
||||
prompt_tokens=prompt_tokens, total_tokens=prompt_tokens
|
||||
)
|
||||
|
||||
@ -0,0 +1,231 @@
|
||||
"""
|
||||
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,
|
||||
transform_openai_input_gemini_embed_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 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."""
|
||||
|
||||
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
|
||||
|
||||
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
|
||||
# Should count tokens only for the text element, not the image
|
||||
assert result.usage.prompt_tokens > 0
|
||||
Loading…
Reference in New Issue
Block a user