diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index 02b69b94d94..11ea4b3619e 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -242,30 +242,41 @@ def _get_embedding_url( ) -> Tuple[str, str]: """ Get URL for embedding models. - + Handles special patterns: - bge/endpoint_id -> strips to endpoint_id for endpoints/ routing - numeric model -> routes to endpoints/ - - regular model -> routes to publishers/google/models/ + - Gemini embedding models -> routes to :embedContent endpoint + - regular model -> routes to publishers/google/models/:predict """ - endpoint = "predict" - + # Import here to avoid circular import + from litellm.llms.vertex_ai.vertex_embeddings.embed_transformation import ( + VertexGeminiEmbeddingConfig, + ) + # Strip routing prefixes (bge/, gemma/, etc.) for endpoint URL construction - model = get_vertex_base_model_name(model=model) - + base_model = get_vertex_base_model_name(model=model) + + # Check if it's a Gemini embedding model that requires :embedContent + # This includes models with "embed/" prefix + if model.startswith("embed/") or VertexGeminiEmbeddingConfig.is_gemini_embedding_model(model): + endpoint = "embedContent" + else: + endpoint = "predict" + # Get base URL (handles global vs regional) base_url = get_vertex_base_url(vertex_location) - - if model.isdigit(): + + if base_model.isdigit(): # https://us-central1-aiplatform.googleapis.com/v1/projects/$PROJECT_ID/locations/us-central1/endpoints/$ENDPOINT_ID:predict # https://aiplatform.googleapis.com/v1/projects/$PROJECT_ID/locations/global/endpoints/$ENDPOINT_ID:predict - url = f"{base_url}/{vertex_api_version}/projects/{vertex_project}/locations/{vertex_location}/endpoints/{model}:{endpoint}" + url = f"{base_url}/{vertex_api_version}/projects/{vertex_project}/locations/{vertex_location}/endpoints/{base_model}:{endpoint}" else: # Regular model -> publisher model - # https://us-central1-aiplatform.googleapis.com/v1/projects/$PROJECT_ID/locations/us-central1/publishers/google/models/{model}:predict - # https://aiplatform.googleapis.com/v1/projects/$PROJECT_ID/locations/global/publishers/google/models/{model}:predict - url = f"{base_url}/v1/projects/{vertex_project}/locations/{vertex_location}/publishers/google/models/{model}:{endpoint}" - + # For Gemini embeddings: https://us-central1-aiplatform.googleapis.com/v1/projects/$PROJECT_ID/locations/us-central1/publishers/google/models/{model}:embedContent + # For other embeddings: https://us-central1-aiplatform.googleapis.com/v1/projects/$PROJECT_ID/locations/us-central1/publishers/google/models/{model}:predict + url = f"{base_url}/v1/projects/{vertex_project}/locations/{vertex_location}/publishers/google/models/{base_model}:{endpoint}" + return url, endpoint diff --git a/litellm/llms/vertex_ai/vertex_embeddings/__init__.py b/litellm/llms/vertex_ai/vertex_embeddings/__init__.py new file mode 100644 index 00000000000..81c45cea326 --- /dev/null +++ b/litellm/llms/vertex_ai/vertex_embeddings/__init__.py @@ -0,0 +1,36 @@ +from .bge import VertexBGEConfig +from .embed_transformation import ( + VertexGeminiEmbeddingConfig as EmbedVertexGeminiEmbeddingConfig, +) +from .transformation import VertexAITextEmbeddingConfig + +__all__ = [ + "VertexAITextEmbeddingConfig", + "VertexBGEConfig", + "EmbedVertexGeminiEmbeddingConfig", + "get_vertex_ai_embedding_config", +] + + +def get_vertex_ai_embedding_config(model: str) -> VertexAITextEmbeddingConfig: + """ + Get the appropriate embedding config for a Vertex AI model. + + Routes to the correct transformation class based on the model type: + - Models with "embed/" prefix use embedContent API (EmbedVertexGeminiEmbeddingConfig) + - BGE models use their own transformation (VertexBGEConfig) + - Other models use predict API (VertexAITextEmbeddingConfig) + + Args: + model: The model name (e.g., "embed/gemini-embedding-2-exp-11-2025", "textembedding-gecko") + + Returns: + VertexAITextEmbeddingConfig: The appropriate configuration class + """ + # Check if model has "embed/" prefix + if model.startswith("embed/"): + return EmbedVertexGeminiEmbeddingConfig() + + # For other models, return the standard config + # The config itself handles routing to Gemini/BGE when needed + return VertexAITextEmbeddingConfig() diff --git a/litellm/llms/vertex_ai/vertex_embeddings/embed_transformation.py b/litellm/llms/vertex_ai/vertex_embeddings/embed_transformation.py new file mode 100644 index 00000000000..8e836aceb2a --- /dev/null +++ b/litellm/llms/vertex_ai/vertex_embeddings/embed_transformation.py @@ -0,0 +1,189 @@ +""" +Gemini Embedding Models Transformation for Vertex AI + +Handles Gemini embedding models that require the :embedContent endpoint +with {"content": {"parts": [{"text": "..."}]}} format instead of the +standard :predict endpoint with {"instances": [...]} format. +""" + +from typing import List, Union + +from litellm.types.llms.vertex_ai import ContentType, PartType +from litellm.types.utils import EmbeddingResponse, Usage + +# List of Gemini embedding models that require :embedContent endpoint +GEMINI_EMBEDDING_MODELS = { + "gemini-embedding-001", + "gemini-embedding-2-exp-11-2025", + "text-embedding-005", + "text-multilingual-embedding-002", +} + + +class VertexGeminiEmbeddingConfig: + """ + Configuration and transformation for Gemini embedding models on Vertex AI. + + These models use the :embedContent endpoint instead of :predict. + """ + + @staticmethod + def is_gemini_embedding_model(model: str) -> bool: + """ + Check if the model is a Gemini embedding model that requires :embedContent endpoint. + + Args: + model: The model name (may include routing prefixes like "vertex_ai/") + + Returns: + bool: True if the model is a Gemini embedding model + """ + # Strip any routing prefixes + base_model = model.split("/")[-1] + return base_model in GEMINI_EMBEDDING_MODELS + + @staticmethod + def transform_openai_request_to_vertex_embedding_request( + input: Union[list, str], optional_params: dict, model: str + ) -> dict: + """Alias for transform_request to match the standard interface.""" + return VertexGeminiEmbeddingConfig.transform_request(input, optional_params, model) + + @staticmethod + def transform_request( + input: Union[list, str], optional_params: dict, model: str + ) -> dict: + """ + Transforms an OpenAI request to a Gemini embedContent request format. + + Gemini embedding models use the :embedContent endpoint with format: + { + "content": {"parts": [{"text": "..."}]}, + "taskType": "...", + "outputDimensionality": ... + } + + Reference: https://cloud.google.com/vertex-ai/generative-ai/docs/embeddings/get-text-embeddings + + Args: + input: Text input(s) to embed + optional_params: Additional parameters (task_type, outputDimensionality, etc.) + model: Model name + + Returns: + dict: Gemini embedContent request format + """ + if isinstance(input, str): + input_list = [input] + else: + input_list = input + + # For single input, use the simple embedContent format + if len(input_list) == 1: + request: dict = {"content": ContentType(parts=[PartType(text=input_list[0])])} + + # Add task type if specified + task_type = optional_params.get("task_type") + if task_type: + request["taskType"] = task_type + + # Add output dimensionality if specified + output_dim = optional_params.get("outputDimensionality") + if output_dim: + request["outputDimensionality"] = output_dim + + return request + else: + # For multiple inputs, we need to call the endpoint multiple times + # Store the full input list in a special key for the handler + request: dict = { + "content": ContentType(parts=[PartType(text=input_list[0])]), + "_batch_inputs": input_list, # Internal flag for handler + } + + # Add task type if specified + task_type = optional_params.get("task_type") + if task_type: + request["taskType"] = task_type + + # Add output dimensionality if specified + output_dim = optional_params.get("outputDimensionality") + if output_dim: + request["outputDimensionality"] = output_dim + + return request + + @staticmethod + def transform_vertex_response_to_openai( + response: Union[dict, List[dict]], model: str, model_response: EmbeddingResponse + ) -> EmbeddingResponse: + """Alias for transform_response to match the standard interface.""" + return VertexGeminiEmbeddingConfig.transform_response(response, model, model_response) + + @staticmethod + def transform_response( + response: Union[dict, List[dict]], model: str, model_response: EmbeddingResponse + ) -> EmbeddingResponse: + """ + Transforms a Gemini embedContent response to OpenAI format. + + Gemini embedContent response format: + { + "embedding": { + "values": [0.1, 0.2, ...] + } + } + + Or for multiple embeddings (from handler looping): + [ + {"embedding": {"values": [...]}, ...}, + {"embedding": {"values": [...]}, ...} + ] + + Args: + response: Gemini embedContent response(s) + model: Model name + model_response: EmbeddingResponse object to populate + + Returns: + EmbeddingResponse: OpenAI-compatible embedding response + """ + embedding_response = [] + + # Check if response is a list (multiple embeddings) or single embedding + if isinstance(response, list): + # Multiple embeddings + for idx, item in enumerate(response): + if "embedding" in item: + embedding_values = item["embedding"]["values"] + embedding_response.append( + { + "object": "embedding", + "index": idx, + "embedding": embedding_values, + } + ) + else: + # Single embedding + if "embedding" in response: + embedding_values = response["embedding"]["values"] + embedding_response.append( + { + "object": "embedding", + "index": 0, + "embedding": embedding_values, + } + ) + + model_response.object = "list" + model_response.data = embedding_response + model_response.model = model + + # Gemini embedContent doesn't return token counts in the response + # Use a basic estimation or set to 0 + input_tokens = 0 + 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/embedding_handler.py b/litellm/llms/vertex_ai/vertex_embeddings/embedding_handler.py index 8a03738ad78..87e4cff2a52 100644 --- a/litellm/llms/vertex_ai/vertex_embeddings/embedding_handler.py +++ b/litellm/llms/vertex_ai/vertex_embeddings/embedding_handler.py @@ -15,6 +15,7 @@ from litellm.llms.vertex_ai.vertex_llm_base import VertexBase from litellm.types.llms.vertex_ai import * from litellm.types.utils import EmbeddingResponse +from . import get_vertex_ai_embedding_config from .types import * @@ -90,8 +91,12 @@ class VertexEmbedding(VertexBase): use_psc_endpoint_format=use_psc_endpoint_format, ) headers = self.set_headers(auth_header=auth_header, extra_headers=extra_headers) + + # Get the appropriate config based on the model + config = get_vertex_ai_embedding_config(model) + vertex_request: VertexEmbeddingRequest = ( - litellm.vertexAITextEmbeddingConfig.transform_openai_request_to_vertex_embedding_request( + config.transform_openai_request_to_vertex_embedding_request( input=input, optional_params=optional_params, model=model ) ) @@ -130,7 +135,7 @@ class VertexEmbedding(VertexBase): ) model_response = ( - litellm.vertexAITextEmbeddingConfig.transform_vertex_response_to_openai( + config.transform_vertex_response_to_openai( response=_json_response, model=model, model_response=model_response ) ) @@ -186,8 +191,12 @@ class VertexEmbedding(VertexBase): use_psc_endpoint_format=use_psc_endpoint_format, ) headers = self.set_headers(auth_header=auth_header, extra_headers=extra_headers) + + # Get the appropriate config based on the model + config = get_vertex_ai_embedding_config(model) + vertex_request: VertexEmbeddingRequest = ( - litellm.vertexAITextEmbeddingConfig.transform_openai_request_to_vertex_embedding_request( + config.transform_openai_request_to_vertex_embedding_request( input=input, optional_params=optional_params, model=model ) ) @@ -228,7 +237,7 @@ class VertexEmbedding(VertexBase): ) model_response = ( - litellm.vertexAITextEmbeddingConfig.transform_vertex_response_to_openai( + config.transform_vertex_response_to_openai( response=_json_response, model=model, model_response=model_response ) ) diff --git a/litellm/llms/vertex_ai/vertex_embeddings/transformation.py b/litellm/llms/vertex_ai/vertex_embeddings/transformation.py index 5a3a4a7188a..99d4de621cf 100644 --- a/litellm/llms/vertex_ai/vertex_embeddings/transformation.py +++ b/litellm/llms/vertex_ai/vertex_embeddings/transformation.py @@ -101,12 +101,23 @@ class VertexAITextEmbeddingConfig(BaseModel): def transform_openai_request_to_vertex_embedding_request( self, input: Union[list, str], optional_params: dict, model: str - ) -> VertexEmbeddingRequest: + ) -> Union[VertexEmbeddingRequest, dict]: """ 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 + from litellm.llms.vertex_ai.vertex_embeddings.embed_transformation import ( + VertexGeminiEmbeddingConfig, + ) + + # Check if it's a Gemini embedding model that requires :embedContent format + # This includes models with "embed/" prefix + if model.startswith("embed/") or VertexGeminiEmbeddingConfig.is_gemini_embedding_model(model): + return VertexGeminiEmbeddingConfig.transform_request( + input=input, optional_params=optional_params, model=model + ) + if model.isdigit(): return self._transform_openai_request_to_fine_tuned_embedding_request( input, optional_params, model @@ -211,14 +222,24 @@ class VertexAITextEmbeddingConfig(BaseModel): """ Transforms a vertex embedding response to an openai response. """ + # Import here to avoid circular import issues with litellm.__init__ + from litellm.llms.vertex_ai.vertex_embeddings.bge import VertexBGEConfig + from litellm.llms.vertex_ai.vertex_embeddings.embed_transformation import ( + VertexGeminiEmbeddingConfig, + ) + + # Check if it's a Gemini embedContent response + # This includes models with "embed/" prefix + if model.startswith("embed/") or VertexGeminiEmbeddingConfig.is_gemini_embedding_model(model): + return VertexGeminiEmbeddingConfig.transform_response( + response=response, model=model, model_response=model_response + ) + if model.isdigit(): return self._transform_vertex_response_to_openai_for_fine_tuned_models( 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 diff --git a/tests/test_litellm/llms/vertex_ai/test_gemini_embeddings.py b/tests/test_litellm/llms/vertex_ai/test_gemini_embeddings.py new file mode 100644 index 00000000000..5330e04b43b --- /dev/null +++ b/tests/test_litellm/llms/vertex_ai/test_gemini_embeddings.py @@ -0,0 +1,155 @@ +""" +Test Gemini Embedding Models on Vertex AI + +Tests the transformation logic for Gemini embedding models that use +the :embedContent endpoint instead of :predict. +""" + +from litellm.llms.vertex_ai.vertex_embeddings.gemini_embeddings import ( + GEMINI_EMBEDDING_MODELS, + VertexGeminiEmbeddingConfig, +) +from litellm.types.utils import EmbeddingResponse + + +class TestVertexGeminiEmbeddingConfig: + """Test VertexGeminiEmbeddingConfig class""" + + def test_is_gemini_embedding_model(self): + """Test model detection for Gemini embedding models""" + # Should detect Gemini embedding models + assert VertexGeminiEmbeddingConfig.is_gemini_embedding_model( + "gemini-embedding-001" + ) + assert VertexGeminiEmbeddingConfig.is_gemini_embedding_model( + "gemini-embedding-2-exp-11-2025" + ) + assert VertexGeminiEmbeddingConfig.is_gemini_embedding_model( + "text-embedding-005" + ) + assert VertexGeminiEmbeddingConfig.is_gemini_embedding_model( + "text-multilingual-embedding-002" + ) + + # Should handle routing prefixes + assert VertexGeminiEmbeddingConfig.is_gemini_embedding_model( + "vertex_ai/gemini-embedding-001" + ) + + # Should not detect non-Gemini models + assert not VertexGeminiEmbeddingConfig.is_gemini_embedding_model( + "textembedding-gecko" + ) + assert not VertexGeminiEmbeddingConfig.is_gemini_embedding_model( + "text-embedding-ada-002" + ) + + def test_transform_request_single_input(self): + """Test request transformation for single input""" + input_text = "Hello, world!" + optional_params = { + "task_type": "RETRIEVAL_QUERY", + "outputDimensionality": 768, + } + + result = VertexGeminiEmbeddingConfig.transform_request( + input=input_text, + optional_params=optional_params, + model="gemini-embedding-001", + ) + + # Verify structure + assert "content" in result + assert "parts" in result["content"] + assert len(result["content"]["parts"]) == 1 + assert result["content"]["parts"][0]["text"] == input_text + + # Verify optional params + assert result["taskType"] == "RETRIEVAL_QUERY" + assert result["outputDimensionality"] == 768 + + # Should not have batch flag for single input + assert "_batch_inputs" not in result + + def test_transform_request_multiple_inputs(self): + """Test request transformation for multiple inputs""" + input_texts = ["Hello, world!", "Goodbye, world!"] + optional_params = {"task_type": "SEMANTIC_SIMILARITY"} + + result = VertexGeminiEmbeddingConfig.transform_request( + input=input_texts, + optional_params=optional_params, + model="gemini-embedding-001", + ) + + # Should have batch flag for multiple inputs + assert "_batch_inputs" in result + assert result["_batch_inputs"] == input_texts + + # First input should be in content + assert "content" in result + assert result["content"]["parts"][0]["text"] == input_texts[0] + + # Verify optional params + assert result["taskType"] == "SEMANTIC_SIMILARITY" + + def test_transform_response_single_embedding(self): + """Test response transformation for single embedding""" + response = {"embedding": {"values": [0.1, 0.2, 0.3]}} + + model_response = EmbeddingResponse() + result = VertexGeminiEmbeddingConfig.transform_response( + response=response, + model="gemini-embedding-001", + model_response=model_response, + ) + + # Verify structure + assert result.object == "list" + assert len(result.data) == 1 + assert result.data[0]["object"] == "embedding" + assert result.data[0]["index"] == 0 + assert result.data[0]["embedding"] == [0.1, 0.2, 0.3] + assert result.model == "gemini-embedding-001" + + # Verify usage + assert hasattr(result, "usage") + assert result.usage.prompt_tokens == 0 # Not provided in response + assert result.usage.total_tokens == 0 + + def test_transform_response_multiple_embeddings(self): + """Test response transformation for multiple embeddings""" + responses = [ + {"embedding": {"values": [0.1, 0.2, 0.3]}}, + {"embedding": {"values": [0.4, 0.5, 0.6]}}, + ] + + model_response = EmbeddingResponse() + result = VertexGeminiEmbeddingConfig.transform_response( + response=responses, + model="gemini-embedding-001", + model_response=model_response, + ) + + # Verify structure + assert result.object == "list" + assert len(result.data) == 2 + + # First embedding + assert result.data[0]["object"] == "embedding" + assert result.data[0]["index"] == 0 + assert result.data[0]["embedding"] == [0.1, 0.2, 0.3] + + # Second embedding + assert result.data[1]["object"] == "embedding" + assert result.data[1]["index"] == 1 + assert result.data[1]["embedding"] == [0.4, 0.5, 0.6] + + assert result.model == "gemini-embedding-001" + + def test_gemini_embedding_models_list(self): + """Test that GEMINI_EMBEDDING_MODELS contains expected models""" + assert "gemini-embedding-001" in GEMINI_EMBEDDING_MODELS + assert "gemini-embedding-2-exp-11-2025" in GEMINI_EMBEDDING_MODELS + assert "text-embedding-005" in GEMINI_EMBEDDING_MODELS + assert "text-multilingual-embedding-002" in GEMINI_EMBEDDING_MODELS