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 834d3a5c0a8..ca1c5d1f58f 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 @@ -313,7 +313,19 @@ def process_response( model_response.model = model if _is_multimodal_input(input): - prompt_tokens = 0 + input_list = input if isinstance(input, list) else [input] + text_elements = [ + e for e in input_list + if isinstance(e, str) + and not (e.startswith("data:") and ";base64," in e) + and not _is_gcs_url(e) + and not _is_file_reference(e) + ] + if text_elements: + input_text = get_formatted_prompt(data={"input": text_elements}, call_type="embedding") + prompt_tokens = token_counter(model=model, text=input_text) + else: + prompt_tokens = 0 else: input_text = get_formatted_prompt(data={"input": input}, call_type="embedding") prompt_tokens = token_counter(model=model, text=input_text) diff --git a/tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py b/tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py index 1417c04503c..c15da3cdee7 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py @@ -20,6 +20,7 @@ from litellm.llms.vertex_ai.gemini_embeddings.batch_embed_content_transformation _is_multimodal_input, process_response, transform_openai_input_gemini_content, + transform_openai_input_gemini_embed_content, ) from litellm.types.llms.vertex_ai import VertexAIBatchEmbeddingsResponseObject from litellm.types.utils import EmbeddingResponse @@ -140,6 +141,40 @@ class TestTransformOpenaiInputGeminiContent: assert len(result["requests"]) == 3 +class TestTransformOpenaiInputGeminiEmbedContent: + """Test transform_openai_input_gemini_embed_content (vertex_ai / embedContent path).""" + + def test_text_and_image_combined(self): + result = transform_openai_input_gemini_embed_content( + input=["hello", IMAGE_DATA_URI], + model="gemini-embedding-2-preview", + optional_params={}, + ) + assert "content" in result + parts = result["content"]["parts"] + assert len(parts) == 2 + assert parts[0]["text"] == "hello" + assert parts[1]["inline_data"] is not None + + def test_gcs_url(self): + result = transform_openai_input_gemini_embed_content( + input=[GCS_URL], + model="gemini-embedding-2-preview", + optional_params={}, + ) + parts = result["content"]["parts"] + assert len(parts) == 1 + assert parts[0]["file_data"]["file_uri"] == GCS_URL + + def test_dimensions_mapped(self): + result = transform_openai_input_gemini_embed_content( + input="hello", + model="gemini-embedding-2-preview", + optional_params={"dimensions": 256}, + ) + assert result["outputDimensionality"] == 256 + + class TestProcessResponse: """Test that process_response sets correct indices.""" @@ -192,4 +227,5 @@ class TestProcessResponse: assert len(result.data) == 2 assert result.data[0]["index"] == 0 assert result.data[1]["index"] == 1 - assert result.usage.prompt_tokens == 0 + # Should count tokens only for the text element, not the image + assert result.usage.prompt_tokens > 0