From 2a9bcf2530e61ef573db9d6955586caa06789eed Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Wed, 11 Mar 2026 11:41:29 +0530 Subject: [PATCH] Fix greptile reviews --- .../batch_embed_content_handler.py | 26 +++---- .../batch_embed_content_transformation.py | 7 +- .../vertex_ai/test_gemini_batch_embeddings.py | 68 +++++++++++++++++++ 3 files changed, 87 insertions(+), 14 deletions(-) diff --git a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py index 1447eb4b92..25c3465807 100644 --- a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py +++ b/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, diff --git a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py index b2bf2c6eb5..41f477d9db 100644 --- a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py +++ b/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 ) diff --git a/tests/litellm/llms/vertex_ai/test_gemini_batch_embeddings.py b/tests/litellm/llms/vertex_ai/test_gemini_batch_embeddings.py index 302facefcb..1ed1de01b5 100644 --- a/tests/litellm/llms/vertex_ai/test_gemini_batch_embeddings.py +++ b/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 +