From 3293ac8a3d282717b32e1b4aa562caede70151eb Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Tue, 28 Oct 2025 18:11:01 -0700 Subject: [PATCH] add BGE handling --- .../llms/vertex_ai/vertex_embeddings/transformation.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/litellm/llms/vertex_ai/vertex_embeddings/transformation.py b/litellm/llms/vertex_ai/vertex_embeddings/transformation.py index caaf00e199e..7bbe13e3597 100644 --- a/litellm/llms/vertex_ai/vertex_embeddings/transformation.py +++ b/litellm/llms/vertex_ai/vertex_embeddings/transformation.py @@ -5,6 +5,7 @@ from pydantic import BaseModel from litellm.types.utils import EmbeddingResponse, Usage +from .bge import VertexBGEConfig from .types import * @@ -109,6 +110,11 @@ class VertexAITextEmbeddingConfig(BaseModel): return self._transform_openai_request_to_fine_tuned_embedding_request( input, optional_params, model ) + + if VertexBGEConfig.is_bge_model(model): + return VertexBGEConfig.transform_request( + input=input, optional_params=optional_params, model=model + ) vertex_request: VertexEmbeddingRequest = VertexEmbeddingRequest() vertex_text_embedding_input_list: List[TextEmbeddingInput] = [] @@ -186,8 +192,8 @@ class VertexAITextEmbeddingConfig(BaseModel): Args: content (str): The content to be embedded. - task_type (Optional[TaskType]): The type of task to be performed". - title (Optional[str]): The title of the document to be embedded + task_type (Optional[TaskType]): The type of task to be performed. + title (Optional[str]): The title of the document to be embedded. Returns: TextEmbeddingInput: A TextEmbeddingInput object.