diff --git a/litellm/llms/vertex_ai/vertex_embeddings/bge.py b/litellm/llms/vertex_ai/vertex_embeddings/bge.py index b8979f55880..2eff0ba96db 100644 --- a/litellm/llms/vertex_ai/vertex_embeddings/bge.py +++ b/litellm/llms/vertex_ai/vertex_embeddings/bge.py @@ -4,6 +4,10 @@ Vertex AI BGE (BAAI General Embedding) Configuration 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 + +Model name handling: +- Model names like "bge/endpoint_id" are automatically transformed in common_utils._get_vertex_url() +- This module focuses on request/response transformation only """ from typing import List, Optional, Union @@ -24,6 +28,13 @@ class VertexBGEConfig: BGE (BAAI General Embedding) models use a different request format where the input field is named "prompt" instead of "content". + + Supported model patterns (after provider split in main.py): + - "bge-small-en-v1.5" (model name) + - "bge/204379420394258432" (endpoint ID pattern) + + Note: Model name transformation (bge/ -> numeric ID) is handled automatically + in common_utils._get_vertex_url(). This class focuses on request/response format only. """ @staticmethod @@ -31,13 +42,19 @@ class VertexBGEConfig: """ Check if the model is a BGE (BAAI General Embedding) model. + After provider split in main.py, supports: + - "bge-small-en-v1.5" (model name) + - "bge/204379420394258432" (endpoint ID pattern) + Args: - model: The model name + model: The model name after provider split Returns: bool: True if the model is a BGE model """ - return "bge" in model.lower() + model_lower = model.lower() + # Check for "bge/" prefix (endpoint pattern) or "bge" in model name + return model_lower.startswith("bge/") or "bge" in model_lower @staticmethod def transform_request( diff --git a/litellm/llms/vertex_ai/vertex_embeddings/transformation.py b/litellm/llms/vertex_ai/vertex_embeddings/transformation.py index 77da3ce7c01..5a3a4a7188a 100644 --- a/litellm/llms/vertex_ai/vertex_embeddings/transformation.py +++ b/litellm/llms/vertex_ai/vertex_embeddings/transformation.py @@ -5,7 +5,6 @@ from pydantic import BaseModel from litellm.types.utils import EmbeddingResponse, Usage -from .bge import VertexBGEConfig from .types import * @@ -106,11 +105,12 @@ class VertexAITextEmbeddingConfig(BaseModel): """ Transforms an openai request to a vertex embedding request. """ + # Import here to avoid circular import issues with litellm.__init__ + from litellm.llms.vertex_ai.vertex_embeddings.bge import VertexBGEConfig if model.isdigit(): 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 @@ -216,6 +216,9 @@ class VertexAITextEmbeddingConfig(BaseModel): response, model, model_response ) + # Import here to avoid circular import issues with litellm.__init__ + from litellm.llms.vertex_ai.vertex_embeddings.bge import VertexBGEConfig + if VertexBGEConfig.is_bge_model(model): return VertexBGEConfig.transform_response( response=response, model=model, model_response=model_response