Fix greptile reviews

litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py

@@ -55,8 +55,9 @@ class GoogleBatchEmbeddings(VertexLLM):
         
         for element in input_list:
             if isinstance(element, str) and _is_file_reference(element):
-                url = f"https://generativelanguage.googleapis.com/v1beta/{element}?key={api_key}"
-                response = sync_handler.get(url=url)
+                url = f"https://generativelanguage.googleapis.com/v1beta/{element}"
+                headers = {"x-goog-api-key": api_key}
+                response = sync_handler.get(url=url, headers=headers)
                 
                 if response.status_code != 200:
                     raise Exception(
@@ -93,8 +94,9 @@ class GoogleBatchEmbeddings(VertexLLM):
         
         for element in input_list:
             if isinstance(element, str) and _is_file_reference(element):
-                url = f"https://generativelanguage.googleapis.com/v1beta/{element}?key={api_key}"
-                response = await async_handler.get(url=url)
+                url = f"https://generativelanguage.googleapis.com/v1beta/{element}"
+                headers = {"x-goog-api-key": api_key}
+                response = await async_handler.get(url=url, headers=headers)
                 
                 if response.status_code != 200:
                     raise Exception(
@@ -151,8 +153,8 @@ class GoogleBatchEmbeddings(VertexLLM):
         optional_params = optional_params or {}
 
         is_multimodal = _is_multimodal_input(input)
-        
-        if is_multimodal:
+        use_embed_content = is_multimodal or (custom_llm_provider == "vertex_ai")
+        if use_embed_content:
             mode = "embedding"
         else:
             mode = "batch_embedding"
@@ -192,14 +194,14 @@ class GoogleBatchEmbeddings(VertexLLM):
                 timeout=timeout,
                 headers=headers,
                 input=input,
-                is_multimodal=is_multimodal,
+                use_embed_content=use_embed_content,
                 api_key=api_key,
                 optional_params=optional_params,
                 logging_obj=logging_obj,
             )
 
         ### TRANSFORMATION (sync path) ###
-        if is_multimodal:
+        if use_embed_content:
             resolved_files = {}
             if api_key:
                 resolved_files = self._resolve_file_references(
@@ -238,7 +240,7 @@ class GoogleBatchEmbeddings(VertexLLM):
 
         _json_response = response.json()
         
-        if is_multimodal:
+        if use_embed_content:
             return process_embed_content_response(
                 input=input,
                 model_response=model_response,
@@ -265,7 +267,7 @@ class GoogleBatchEmbeddings(VertexLLM):
         timeout: Optional[Union[float, httpx.Timeout]],
         headers={},
         client: Optional[AsyncHTTPHandler] = None,
-        is_multimodal: bool = False,
+        use_embed_content: bool = False,
         api_key: Optional[str] = None,
         optional_params: Optional[dict] = None,
         logging_obj: Optional[Any] = None,
@@ -287,7 +289,7 @@ class GoogleBatchEmbeddings(VertexLLM):
             async_handler = client  # type: ignore
 
         ### TRANSFORMATION (async path) ###
-        if is_multimodal:
+        if use_embed_content:
             resolved_files = {}
             if api_key:
                 resolved_files = await self._async_resolve_file_references(
@@ -327,7 +329,7 @@ class GoogleBatchEmbeddings(VertexLLM):
 
         _json_response = response.json()
         
-        if is_multimodal:
+        if use_embed_content:
             return process_embed_content_response(
                 input=input,
                 model_response=model_response,

litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py

@@ -267,8 +267,11 @@ def process_embed_content_response(
     model_response.data = [openai_embedding]
     model_response.model = model
     
-    input_text = get_formatted_prompt(data={"input": input}, call_type="embedding")
-    prompt_tokens = token_counter(model=model, text=input_text)
+    if _is_multimodal_input(input):
+        prompt_tokens = 0
+    else:
+        input_text = get_formatted_prompt(data={"input": input}, call_type="embedding")
+        prompt_tokens = token_counter(model=model, text=input_text)
     model_response.usage = Usage(
         prompt_tokens=prompt_tokens, total_tokens=prompt_tokens
     )

tests/litellm/llms/vertex_ai/test_gemini_batch_embeddings.py

@@ -271,6 +271,26 @@ def test_embed_content_response_processing():
     assert result.data[0].index == 0
     assert result.data[0].object == "embedding"
     assert result.model == "gemini-embedding-2-preview"
+    assert result.usage.prompt_tokens > 0
+
+
+def test_embed_content_response_multimodal_sets_prompt_tokens_zero():
+    """Test that multimodal input sets prompt_tokens=0 (cannot accurately count)."""
+    response_json = {
+        "embedding": {
+            "values": [0.1, 0.2, 0.3, 0.4, 0.5]
+        }
+    }
+    
+    model_response = EmbeddingResponse()
+    result = process_embed_content_response(
+        input=["text", "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg=="],
+        model_response=model_response,
+        model="gemini-embedding-2-preview",
+        response_json=response_json,
+    )
+    
+    assert result.usage.prompt_tokens == 0
 
 
 def test_gemini_multimodal_embedding_e2e():
@@ -456,3 +476,51 @@ def test_multimodal_input_detection_with_gcs():
     assert _is_multimodal_input("gs://bucket/video.mp4") is True
     assert _is_multimodal_input(["just text", "more text"]) is False
 
+
+def test_vertex_ai_text_only_embedding_uses_embed_content():
+    """
+    Test that vertex_ai/gemini-embedding-2-preview with text-only input uses
+    embedContent endpoint (not batchEmbedContents) and returns a single embedding.
+    """
+    client = HTTPHandler()
+    embed_content_url = "https://us-central1-aiplatform.googleapis.com/v1/projects/test/locations/us-central1/publishers/google/models/gemini-embedding-2-preview:embedContent"
+
+    def mock_auth_token(*args, **kwargs):
+        return "Bearer test-token", "test-project"
+
+    with patch.object(client, "post") as mock_post, patch(
+        "litellm.llms.vertex_ai.gemini_embeddings.batch_embed_content_handler.GoogleBatchEmbeddings._ensure_access_token",
+        side_effect=mock_auth_token,
+    ), patch(
+        "litellm.llms.vertex_ai.gemini_embeddings.batch_embed_content_handler.GoogleBatchEmbeddings._get_token_and_url"
+    ) as mock_get_token:
+        mock_get_token.return_value = (
+            {"Authorization": "Bearer test-token"},
+            embed_content_url,
+        )
+        mock_response = MagicMock()
+        mock_response.status_code = 200
+        mock_response.json.return_value = {
+            "embedding": {"values": [0.1, 0.2, 0.3, 0.4, 0.5]}
+        }
+        mock_post.return_value = mock_response
+
+        response = litellm.embedding(
+            model="vertex_ai/gemini-embedding-2-preview",
+            input=["Hello, world!"],
+            vertex_project="test-project",
+            vertex_location="us-central1",
+            client=client,
+        )
+
+        mock_post.assert_called_once()
+        call_args = mock_post.call_args
+        post_url = call_args.kwargs.get("url", call_args.args[0] if call_args.args else "")
+        assert "embedContent" in str(post_url)
+        data = json.loads(call_args.kwargs["data"])
+        assert "content" in data
+        assert "parts" in data["content"]
+        assert len(data["content"]["parts"]) == 1
+        assert data["content"]["parts"][0]["text"] == "Hello, world!"
+        assert len(response.data) == 1
+