diff --git a/litellm/llms/vertex_ai/multimodal_embeddings/transformation.py b/litellm/llms/vertex_ai/multimodal_embeddings/transformation.py index a4815b01ecc..90a2e4bf653 100644 --- a/litellm/llms/vertex_ai/multimodal_embeddings/transformation.py +++ b/litellm/llms/vertex_ai/multimodal_embeddings/transformation.py @@ -287,14 +287,14 @@ class VertexAIMultimodalEmbeddingConfig(BaseEmbeddingConfig): object="embedding", ) openai_embeddings.append(openai_embedding_object) - elif "imageEmbedding" in _prediction: + if "imageEmbedding" in _prediction: openai_embedding_object = Embedding( embedding=_prediction["imageEmbedding"], index=idx, object="embedding", ) openai_embeddings.append(openai_embedding_object) - elif "videoEmbeddings" in _prediction: + if "videoEmbeddings" in _prediction: for video_embedding in _prediction["videoEmbeddings"]: openai_embedding_object = Embedding( embedding=video_embedding["embedding"], diff --git a/tests/local_testing/test_amazing_vertex_completion.py b/tests/local_testing/test_amazing_vertex_completion.py index 6e31166ad99..1716b3d8ae7 100644 --- a/tests/local_testing/test_amazing_vertex_completion.py +++ b/tests/local_testing/test_amazing_vertex_completion.py @@ -1880,8 +1880,9 @@ async def test_vertexai_multimodal_embedding(): assert args_to_vertexai == expected_payload assert response.model == "multimodalembedding@001" - assert len(response.data) == 1 - response_data = response.data[0] + assert len(response.data) == 2 + assert response.data[0]["embedding"] == [0.4, 0.5, 0.6] + assert response.data[1]["embedding"] == [0.1, 0.2, 0.3] # Optional: Print for debugging print("Arguments passed to Vertex AI:", args_to_vertexai) diff --git a/tests/test_litellm/llms/vertex_ai/multimodal_embeddings/test_vertex_ai_multimodal_embedding_transformation.py b/tests/test_litellm/llms/vertex_ai/multimodal_embeddings/test_vertex_ai_multimodal_embedding_transformation.py index 6b605aed0ca..676fea9b476 100644 --- a/tests/test_litellm/llms/vertex_ai/multimodal_embeddings/test_vertex_ai_multimodal_embedding_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/multimodal_embeddings/test_vertex_ai_multimodal_embedding_transformation.py @@ -179,3 +179,49 @@ class TestVertexMultimodalEmbedding: headers={}, ) assert "parameters" not in request + + def test_response_keeps_both_text_and_image_embeddings(self): + predictions = { + "predictions": [ + { + "textEmbedding": [0.4, 0.5, 0.6], + "imageEmbedding": [0.1, 0.2, 0.3], + } + ] + } + result = self.config.transform_embedding_response_to_openai(predictions) + assert [e["embedding"] for e in result] == [[0.4, 0.5, 0.6], [0.1, 0.2, 0.3]] + assert [e["index"] for e in result] == [0, 0] + + def test_response_keeps_text_and_video_embeddings(self): + predictions = { + "predictions": [ + { + "textEmbedding": [0.4, 0.5, 0.6], + "videoEmbeddings": [ + { + "startOffsetSec": 0, + "endOffsetSec": 5, + "embedding": [0.7, 0.8, 0.9], + } + ], + } + ] + } + result = self.config.transform_embedding_response_to_openai(predictions) + assert [e["embedding"] for e in result] == [[0.4, 0.5, 0.6], [0.7, 0.8, 0.9]] + + @pytest.mark.parametrize( + "prediction, expected", + [ + ({"textEmbedding": [0.1, 0.2]}, [[0.1, 0.2]]), + ({"imageEmbedding": [0.3, 0.4]}, [[0.3, 0.4]]), + ( + {"videoEmbeddings": [{"startOffsetSec": 0, "endOffsetSec": 5, "embedding": [0.5]}]}, + [[0.5]], + ), + ], + ) + def test_response_single_modality_unchanged(self, prediction, expected): + result = self.config.transform_embedding_response_to_openai({"predictions": [prediction]}) + assert [e["embedding"] for e in result] == expected