mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-11 22:51:28 +00:00
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:
parent
e64d98f725
commit
dd93d2698b
2 changed files with 32 additions and 2 deletions
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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}"
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue