From 2c4a495619fd95cf203bb9b60485ebf6fb2a2377 Mon Sep 17 00:00:00 2001
From: Sameer Kankute <sameer@berri.ai>
Date: Wed, 11 Mar 2026 10:39:34 +0530
Subject: [PATCH] Add support for vertex ai gemini multimodal embedings

---
 litellm/llms/vertex_ai/common_utils.py | 24 ++++++++++++---------
 litellm/main.py                        | 29 ++++++++++++++++++++++++--
 litellm/types/utils.py                 |  1 +
 litellm/utils.py                       |  1 +
 4 files changed, 43 insertions(+), 12 deletions(-)

diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py
index 3c5cbb6543..c02d63414c 100644
--- a/litellm/llms/vertex_ai/common_utils.py
+++ b/litellm/llms/vertex_ai/common_utils.py
@@ -247,23 +247,27 @@ def _get_embedding_url(
     - bge/endpoint_id -> strips to endpoint_id for endpoints/ routing
     - numeric model -> routes to endpoints/
     - regular model -> routes to publishers/google/models/
-    """
-    endpoint = "predict"
-    
-    # Strip routing prefixes (bge/, gemma/, etc.) for endpoint URL construction
+    - models with uses_embed_content flag -> use embedContent endpoint instead of predict
+    """    
+    original_model = model
     model = get_vertex_base_model_name(model=model)
     
-    # Get base URL (handles global vs regional)
+    try:
+        model_info = litellm.get_model_info(
+            model=original_model,
+            custom_llm_provider="vertex_ai",
+        )
+        uses_embed_content = model_info.get("uses_embed_content", False)
+    except Exception:
+        uses_embed_content = False
+    
+    endpoint = "embedContent" if uses_embed_content else "predict"
+    
     base_url = get_vertex_base_url(vertex_location)
     
     if 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}"
     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}"
     
     return url, endpoint
diff --git a/litellm/main.py b/litellm/main.py
index 364519e1fe..529b998810 100644
--- a/litellm/main.py
+++ b/litellm/main.py
@@ -132,6 +132,7 @@ from litellm.utils import (
     create_tokenizer,
     get_api_key,
     get_llm_provider,
+    get_model_info,
     get_non_default_completion_params,
     get_non_default_transcription_params,
     get_optional_params_embeddings,
@@ -5190,13 +5191,37 @@ def embedding(  # noqa: PLR0915
                 or get_secret_str("VERTEX_API_BASE")
             )
 
-            if (
+            try:
+                model_info = get_model_info(model=model, custom_llm_provider="vertex_ai")
+                uses_embed_content = model_info.get("uses_embed_content", False)
+            except Exception:
+                uses_embed_content = False
+
+            if uses_embed_content:
+                response = google_batch_embeddings.batch_embeddings(  # type: ignore
+                    model=model,
+                    input=input,
+                    encoding=_get_encoding(),
+                    logging_obj=logging,
+                    optional_params=optional_params,
+                    model_response=EmbeddingResponse(),
+                    vertex_project=vertex_ai_project,
+                    vertex_location=vertex_ai_location,
+                    vertex_credentials=vertex_credentials,
+                    aembedding=aembedding,
+                    print_verbose=print_verbose,
+                    custom_llm_provider="vertex_ai",
+                    api_key=None,
+                    api_base=api_base,
+                    client=client,
+                    extra_headers=headers,
+                )
+            elif (
                 "image" in optional_params
                 or "video" in optional_params
                 or model
                 in vertex_multimodal_embedding.SUPPORTED_MULTIMODAL_EMBEDDING_MODELS
             ):
-                # multimodal embedding is supported on vertex httpx
                 response = vertex_multimodal_embedding.multimodal_embedding(
                     model=model,
                     input=input,
diff --git a/litellm/types/utils.py b/litellm/types/utils.py
index b5d5c06924..52221c47de 100644
--- a/litellm/types/utils.py
+++ b/litellm/types/utils.py
@@ -253,6 +253,7 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False):
     tpm: Optional[int]
     rpm: Optional[int]
     provider_specific_entry: Optional[Dict[str, float]]
+    uses_embed_content: Optional[bool]
 
 
 class ModelInfo(ModelInfoBase, total=False):
diff --git a/litellm/utils.py b/litellm/utils.py
index 4367ec789b..eebcdb3ec5 100644
--- a/litellm/utils.py
+++ b/litellm/utils.py
@@ -5779,6 +5779,7 @@ def _get_model_info_helper(  # noqa: PLR0915
                 provider_specific_entry=_model_info.get(
                     "provider_specific_entry", None
                 ),
+                uses_embed_content=_model_info.get("uses_embed_content", None),
             )
     except Exception as e:
         verbose_logger.debug(f"Error getting model info: {e}")