fix request transform vertex BGE

This commit is contained in:
Ishaan Jaffer 2025-10-28 18:14:25 -07:00 committed by Sameer Kankute
parent f2befcf657
commit c821acd61a
2 changed files with 57 additions and 2 deletions

View file

@ -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

View file

@ -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"]