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 08831a8215f..389a3a85f56 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 @@ -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) diff --git a/tests/litellm/llms/vertex_ai/test_gemini_batch_embeddings.py b/tests/litellm/llms/vertex_ai/test_gemini_batch_embeddings.py index a8e427d3bc1..d814f8ec97f 100644 --- a/tests/litellm/llms/vertex_ai/test_gemini_batch_embeddings.py +++ b/tests/litellm/llms/vertex_ai/test_gemini_batch_embeddings.py @@ -22,6 +22,7 @@ _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 mock_auth_token(*args, **kwargs): 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}" +