fix(gemini): assign correct indices in batch embedding response (#25656)

### Background

The Gemini batchEmbedContents response handler hardcoded `index=0` for
every embedding in the response. Any consumer relying on the OpenAI-format
`index` field to match embeddings back to inputs would silently get wrong
associations.

### Changes

Use `enumerate` in `process_response` so each embedding gets its
positional index instead of 0.

### Test Plan

Added unit test asserting sequential indices and correct vector ordering
for a 3-element batch response.
This commit is contained in:
lucassz 2026-04-13 18:37:41 -07:00 committed by Sameer Kankute
parent e64d98f725
commit dd93d2698b
No known key found for this signature in database
2 changed files with 32 additions and 2 deletions

View file

@ -292,10 +292,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)

View file

@ -22,6 +22,7 @@ from litellm.llms.vertex_ai.gemini_embeddings.batch_embed_content_transformation
_is_multimodal_input,
_parse_data_url,
process_embed_content_response,
process_response,
transform_openai_input_gemini_content,
transform_openai_input_gemini_embed_content,
)
@ -563,3 +564,32 @@ def test_vertex_ai_text_only_embedding_uses_embed_content():
assert data["content"]["parts"][0]["text"] == "Hello, world!"
assert len(response.data) == 1
def test_batch_embeddings_response_has_correct_indices_and_order():
"""Test that process_response assigns sequential indices and preserves order."""
response_json = {
"embeddings": [
{"values": [0.1, 0.2, 0.3]},
{"values": [0.4, 0.5, 0.6]},
{"values": [0.7, 0.8, 0.9]},
]
}
expected_values = [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6], [0.7, 0.8, 0.9]]
model_response = EmbeddingResponse()
result = process_response(
input=["first", "second", "third"],
model_response=model_response,
model="text-embedding-004",
_predictions=response_json,
)
assert len(result.data) == 3
for i, embedding in enumerate(result.data):
assert (
embedding.index == i
), f"embedding {i} has index={embedding.index}, expected {i}"
assert (
embedding.embedding == expected_values[i]
), f"embedding {i} has wrong values: {embedding.embedding}"