Add embedcontent

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
 
 

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()

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

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

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

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