From c821acd61a1eef41589df3410b8d2d23d9395ab0 Mon Sep 17 00:00:00 2001
From: Ishaan Jaffer <ishaanjaffer0324@gmail.com>
Date: Tue, 28 Oct 2025 18:14:25 -0700
Subject: [PATCH] fix request transform vertex BGE

---
 .../llms/vertex_ai/vertex_embeddings/bge.py   | 54 ++++++++++++++++++-
 .../vertex_embeddings/transformation.py       |  5 ++
 2 files changed, 57 insertions(+), 2 deletions(-)

diff --git a/litellm/llms/vertex_ai/vertex_embeddings/bge.py b/litellm/llms/vertex_ai/vertex_embeddings/bge.py
index 401f7ebd907..1bfa362ee98 100644
--- a/litellm/llms/vertex_ai/vertex_embeddings/bge.py
+++ b/litellm/llms/vertex_ai/vertex_embeddings/bge.py
@@ -1,12 +1,15 @@
 """
 Vertex AI BGE (BAAI General Embedding) Configuration
 
-BGE models deployed on Vertex AI require different input format:
-- Use "prompt" instead of "content" as the input field
+BGE models deployed on Vertex AI require different input/output format:
+- Request: Use "prompt" instead of "content" as the input field
+- Response: Embeddings are returned directly as arrays, not wrapped in objects
 """
 
 from typing import List, Optional, Union
 
+from litellm.types.utils import EmbeddingResponse, Usage
+
 from .types import (
     EmbeddingParameters,
     TaskType,
@@ -98,3 +101,50 @@ class VertexBGEConfig:
             text_embedding_input["title"] = title
         return text_embedding_input
 
+    @staticmethod
+    def transform_response(
+        response: dict, model: str, model_response: EmbeddingResponse
+    ) -> EmbeddingResponse:
+        """
+        Transforms a Vertex BGE embedding response to OpenAI format.
+        
+        BGE models return embeddings directly as arrays in predictions:
+        {
+          "predictions": [
+            [0.002, 0.021, ...],
+            [0.003, 0.022, ...]
+          ]
+        }
+        
+        Args:
+            response: The raw response from Vertex AI
+            model: The model name
+            model_response: The EmbeddingResponse object to populate
+            
+        Returns:
+            EmbeddingResponse: The transformed response in OpenAI format
+        """
+        _predictions = response["predictions"]
+
+        embedding_response = []
+        # BGE models don't return token counts, so we estimate or set to 0
+        input_tokens = 0
+
+        for idx, embedding_values in enumerate(_predictions):
+            embedding_response.append(
+                {
+                    "object": "embedding",
+                    "index": idx,
+                    "embedding": embedding_values,
+                }
+            )
+
+        model_response.object = "list"
+        model_response.data = embedding_response
+        model_response.model = model
+        usage = Usage(
+            prompt_tokens=input_tokens, completion_tokens=0, total_tokens=input_tokens
+        )
+        setattr(model_response, "usage", usage)
+        return model_response
+
diff --git a/litellm/llms/vertex_ai/vertex_embeddings/transformation.py b/litellm/llms/vertex_ai/vertex_embeddings/transformation.py
index 7bbe13e3597..77da3ce7c01 100644
--- a/litellm/llms/vertex_ai/vertex_embeddings/transformation.py
+++ b/litellm/llms/vertex_ai/vertex_embeddings/transformation.py
@@ -215,6 +215,11 @@ class VertexAITextEmbeddingConfig(BaseModel):
             return self._transform_vertex_response_to_openai_for_fine_tuned_models(
                 response, model, model_response
             )
+        
+        if VertexBGEConfig.is_bge_model(model):
+            return VertexBGEConfig.transform_response(
+                response=response, model=model, model_response=model_response
+            )
 
         _predictions = response["predictions"]