From b108c02fd780abd2247bd248340afaf43b3c5177 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Wed, 11 Mar 2026 10:08:17 +0530 Subject: [PATCH 1/7] Add support for gemini multimodal embedings --- .../batch_embed_content_handler.py | 253 +++++++++++++---- .../batch_embed_content_transformation.py | 191 ++++++++++++- litellm/types/llms/vertex_ai.py | 11 + .../vertex_ai/test_gemini_batch_embeddings.py | 264 +++++++++++++++++- 4 files changed, 656 insertions(+), 63 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 07f57a4a7f6..1447eb4b92a 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 @@ -3,12 +3,11 @@ Google AI Studio /batchEmbedContents Embeddings Endpoint """ import json -from typing import Any, Literal, Optional, Union +from typing import Any, Dict, List, Literal, Optional, Union import httpx import litellm -from litellm.types.utils import EmbeddingResponse from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, HTTPHandler, @@ -16,18 +15,100 @@ from litellm.llms.custom_httpx.http_handler import ( ) from litellm.types.llms.openai import EmbeddingInput from litellm.types.llms.vertex_ai import ( + GeminiEmbedContentResponseObject, VertexAIBatchEmbeddingsRequestBody, VertexAIBatchEmbeddingsResponseObject, ) +from litellm.types.utils import EmbeddingResponse from ..gemini.vertex_and_google_ai_studio_gemini import VertexLLM from .batch_embed_content_transformation import ( + _is_file_reference, + _is_multimodal_input, + process_embed_content_response, process_response, transform_openai_input_gemini_content, + transform_openai_input_gemini_embed_content, ) class GoogleBatchEmbeddings(VertexLLM): + def _resolve_file_references( + self, + input: EmbeddingInput, + api_key: str, + sync_handler: HTTPHandler, + ) -> Dict[str, Dict[str, str]]: + """ + Resolve Gemini file references (files/...) to get mime_type and uri. + + Args: + input: EmbeddingInput that may contain file references + api_key: Gemini API key + sync_handler: HTTP client + + Returns: + Dict mapping file name to {mime_type, uri} + """ + input_list = [input] if isinstance(input, str) else input + resolved_files: Dict[str, Dict[str, str]] = {} + + 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) + + if response.status_code != 200: + raise Exception( + f"Error fetching file {element}: {response.status_code} {response.text}" + ) + + file_data = response.json() + resolved_files[element] = { + "mime_type": file_data.get("mimeType", ""), + "uri": file_data.get("uri", element), + } + + return resolved_files + + async def _async_resolve_file_references( + self, + input: EmbeddingInput, + api_key: str, + async_handler: AsyncHTTPHandler, + ) -> Dict[str, Dict[str, str]]: + """ + Async version of _resolve_file_references. + + Args: + input: EmbeddingInput that may contain file references + api_key: Gemini API key + async_handler: Async HTTP client + + Returns: + Dict mapping file name to {mime_type, uri} + """ + input_list = [input] if isinstance(input, str) else input + resolved_files: Dict[str, Dict[str, str]] = {} + + 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) + + if response.status_code != 200: + raise Exception( + f"Error fetching file {element}: {response.status_code} {response.text}" + ) + + file_data = response.json() + resolved_files[element] = { + "mime_type": file_data.get("mimeType", ""), + "uri": file_data.get("uri", element), + } + + return resolved_files + def batch_embeddings( self, model: str, @@ -54,20 +135,6 @@ class GoogleBatchEmbeddings(VertexLLM): custom_llm_provider=custom_llm_provider, ) - auth_header, url = self._get_token_and_url( - model=model, - auth_header=_auth_header, - gemini_api_key=api_key, - vertex_project=vertex_project, - vertex_location=vertex_location, - vertex_credentials=vertex_credentials, - stream=None, - custom_llm_provider=custom_llm_provider, - api_base=api_base, - should_use_v1beta1_features=False, - mode="batch_embedding", - ) - if client is None: _params = {} if timeout is not None: @@ -83,9 +150,25 @@ class GoogleBatchEmbeddings(VertexLLM): optional_params = optional_params or {} - ### TRANSFORMATION ### - request_data = transform_openai_input_gemini_content( - input=input, model=model, optional_params=optional_params + is_multimodal = _is_multimodal_input(input) + + if is_multimodal: + mode = "embedding" + else: + mode = "batch_embedding" + + auth_header, url = self._get_token_and_url( + model=model, + auth_header=_auth_header, + gemini_api_key=api_key, + vertex_project=vertex_project, + vertex_location=vertex_location, + vertex_credentials=vertex_credentials, + stream=None, + custom_llm_provider=custom_llm_provider, + api_base=api_base, + should_use_v1beta1_features=False, + mode=mode, ) headers = { @@ -93,14 +176,46 @@ class GoogleBatchEmbeddings(VertexLLM): } if auth_header is not None: if isinstance(auth_header, dict): - # For Gemini with custom api_base: auth_header is {"x-goog-api-key": "..."} headers.update(auth_header) else: - # For Vertex AI: auth_header is a Bearer token string headers["Authorization"] = f"Bearer {auth_header}" if extra_headers is not None: headers.update(extra_headers) + if aembedding is True: + return self.async_batch_embeddings( # type: ignore + model=model, + api_base=api_base, + url=url, + data=None, + model_response=model_response, + timeout=timeout, + headers=headers, + input=input, + is_multimodal=is_multimodal, + api_key=api_key, + optional_params=optional_params, + logging_obj=logging_obj, + ) + + ### TRANSFORMATION (sync path) ### + if is_multimodal: + resolved_files = {} + if api_key: + resolved_files = self._resolve_file_references( + input=input, api_key=api_key, sync_handler=sync_handler + ) + request_data = transform_openai_input_gemini_embed_content( + input=input, + model=model, + optional_params=optional_params, + resolved_files=resolved_files, + ) + else: + request_data = transform_openai_input_gemini_content( + input=input, model=model, optional_params=optional_params + ) + ## LOGGING logging_obj.pre_call( input=input, @@ -112,18 +227,6 @@ class GoogleBatchEmbeddings(VertexLLM): }, ) - if aembedding is True: - return self.async_batch_embeddings( # type: ignore - model=model, - api_base=api_base, - url=url, - data=request_data, - model_response=model_response, - timeout=timeout, - headers=headers, - input=input, - ) - response = sync_handler.post( url=url, headers=headers, @@ -134,26 +237,38 @@ class GoogleBatchEmbeddings(VertexLLM): raise Exception(f"Error: {response.status_code} {response.text}") _json_response = response.json() - _predictions = VertexAIBatchEmbeddingsResponseObject(**_json_response) # type: ignore - - return process_response( - model=model, - model_response=model_response, - _predictions=_predictions, - input=input, - ) + + if is_multimodal: + return process_embed_content_response( + input=input, + model_response=model_response, + model=model, + response_json=_json_response, + ) + else: + _predictions = VertexAIBatchEmbeddingsResponseObject(**_json_response) # type: ignore + return process_response( + model=model, + model_response=model_response, + _predictions=_predictions, + input=input, + ) async def async_batch_embeddings( self, model: str, api_base: Optional[str], url: str, - data: VertexAIBatchEmbeddingsRequestBody, + data: Optional[Union[VertexAIBatchEmbeddingsRequestBody, dict]], model_response: EmbeddingResponse, input: EmbeddingInput, timeout: Optional[Union[float, httpx.Timeout]], headers={}, client: Optional[AsyncHTTPHandler] = None, + is_multimodal: bool = False, + api_key: Optional[str] = None, + optional_params: Optional[dict] = None, + logging_obj: Optional[Any] = None, ) -> EmbeddingResponse: if client is None: _params = {} @@ -171,6 +286,36 @@ class GoogleBatchEmbeddings(VertexLLM): else: async_handler = client # type: ignore + ### TRANSFORMATION (async path) ### + if is_multimodal: + resolved_files = {} + if api_key: + resolved_files = await self._async_resolve_file_references( + input=input, api_key=api_key, async_handler=async_handler + ) + data = transform_openai_input_gemini_embed_content( + input=input, + model=model, + optional_params=optional_params or {}, + resolved_files=resolved_files, + ) + else: + data = transform_openai_input_gemini_content( + input=input, model=model, optional_params=optional_params or {} + ) + + ## LOGGING + if logging_obj is not None: + logging_obj.pre_call( + input=input, + api_key="", + additional_args={ + "complete_input_dict": data, + "api_base": url, + "headers": headers, + }, + ) + response = await async_handler.post( url=url, headers=headers, @@ -181,11 +326,19 @@ class GoogleBatchEmbeddings(VertexLLM): raise Exception(f"Error: {response.status_code} {response.text}") _json_response = response.json() - _predictions = VertexAIBatchEmbeddingsResponseObject(**_json_response) # type: ignore - - return process_response( - model=model, - model_response=model_response, - _predictions=_predictions, - input=input, - ) + + if is_multimodal: + return process_embed_content_response( + input=input, + model_response=model_response, + model=model, + response_json=_json_response, + ) + else: + _predictions = VertexAIBatchEmbeddingsResponseObject(**_json_response) # type: ignore + return process_response( + model=model, + model_response=model_response, + _predictions=_predictions, + input=input, + ) 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 455ec1d18f5..6070c70677b 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 @@ -4,20 +4,100 @@ Transformation logic from OpenAI /v1/embeddings format to Google AI Studio /batc Why separate file? Make it easy to see how transformation works """ -from typing import List +from typing import Dict, List, Optional, Tuple -from litellm.types.utils import EmbeddingResponse from litellm.types.llms.openai import EmbeddingInput from litellm.types.llms.vertex_ai import ( + BlobType, ContentType, EmbedContentRequest, + FileDataType, PartType, VertexAIBatchEmbeddingsRequestBody, VertexAIBatchEmbeddingsResponseObject, ) -from litellm.types.utils import Embedding, Usage +from litellm.types.utils import Embedding, EmbeddingResponse, Usage from litellm.utils import get_formatted_prompt, token_counter +SUPPORTED_EMBEDDING_MIME_TYPES = { + "image/png", + "image/jpeg", + "audio/mpeg", + "audio/wav", + "video/mp4", + "video/quicktime", + "application/pdf", +} + + +def _is_file_reference(s: str) -> bool: + """Check if string is a Gemini file reference (files/...).""" + return isinstance(s, str) and s.startswith("files/") + + +def _parse_data_url(data_url: str) -> Tuple[str, str]: + """ + Parse a data URL to extract the media type and base64 data. + + Args: + data_url: Data URL in format: data:image/jpeg;base64,/9j/4AAQ... + + Returns: + tuple: (media_type, base64_data) + media_type: e.g., "image/jpeg", "video/mp4", "audio/mpeg" + base64_data: The base64-encoded data without the prefix + + Raises: + ValueError: If data URL format is invalid or MIME type is unsupported + """ + if not data_url.startswith("data:"): + raise ValueError(f"Invalid data URL format: {data_url[:50]}...") + + if "," not in data_url: + raise ValueError(f"Invalid data URL format (missing comma): {data_url[:50]}...") + + metadata, base64_data = data_url.split(",", 1) + + metadata = metadata[5:] + + if ";" in metadata: + media_type = metadata.split(";")[0] + else: + media_type = metadata + + if media_type not in SUPPORTED_EMBEDDING_MIME_TYPES: + raise ValueError( + f"Unsupported MIME type for embedding: {media_type}. " + f"Supported types: {', '.join(sorted(SUPPORTED_EMBEDDING_MIME_TYPES))}" + ) + + return media_type, base64_data + + +def _is_multimodal_input(input: EmbeddingInput) -> bool: + """ + Check if the input contains multimodal data (data URIs or file references). + + Args: + input: EmbeddingInput (str or List[str]) + + Returns: + bool: True if any element is a data URI or file reference + """ + if isinstance(input, str): + input_list = [input] + else: + input_list = input + + for element in input_list: + if isinstance(element, str): + if element.startswith("data:") and ";base64," in element: + return True + if _is_file_reference(element): + return True + + return False + def transform_openai_input_gemini_content( input: EmbeddingInput, model: str, optional_params: dict @@ -26,12 +106,17 @@ def transform_openai_input_gemini_content( The content to embed. Only the parts.text fields will be counted. """ gemini_model_name = "models/{}".format(model) + + gemini_params = optional_params.copy() + if "dimensions" in gemini_params: + gemini_params["outputDimensionality"] = gemini_params.pop("dimensions") + requests: List[EmbedContentRequest] = [] if isinstance(input, str): request = EmbedContentRequest( model=gemini_model_name, content=ContentType(parts=[PartType(text=input)]), - **optional_params + **gemini_params ) requests.append(request) else: @@ -39,13 +124,109 @@ def transform_openai_input_gemini_content( request = EmbedContentRequest( model=gemini_model_name, content=ContentType(parts=[PartType(text=i)]), - **optional_params + **gemini_params ) requests.append(request) return VertexAIBatchEmbeddingsRequestBody(requests=requests) +def transform_openai_input_gemini_embed_content( + input: EmbeddingInput, + model: str, + optional_params: dict, + resolved_files: Optional[Dict[str, Dict[str, str]]] = None, +) -> dict: + """ + Transform OpenAI embedding input to Gemini embedContent format (multimodal). + + Args: + input: EmbeddingInput (str or List[str]) with text, data URIs, or file references + model: Model name + optional_params: Additional parameters (taskType, outputDimensionality, etc.) + resolved_files: Dict mapping file names (files/abc) to {mime_type, uri} + + Returns: + dict: Gemini embedContent request body with content.parts + """ + resolved_files = resolved_files or {} + + gemini_params = optional_params.copy() + if "dimensions" in gemini_params: + gemini_params["outputDimensionality"] = gemini_params.pop("dimensions") + + input_list = [input] if isinstance(input, str) else input + parts: List[PartType] = [] + + for element in input_list: + if not isinstance(element, str): + raise ValueError(f"Unsupported input type: {type(element)}") + + if element.startswith("data:") and ";base64," in element: + mime_type, base64_data = _parse_data_url(element) + blob: BlobType = {"mime_type": mime_type, "data": base64_data} + parts.append(PartType(inline_data=blob)) + elif _is_file_reference(element): + if element not in resolved_files: + raise ValueError(f"File reference {element} not resolved") + file_info = resolved_files[element] + file_data: FileDataType = { + "mime_type": file_info["mime_type"], + "file_uri": file_info["uri"], + } + parts.append(PartType(file_data=file_data)) + else: + parts.append(PartType(text=element)) + + request_body: dict = { + "content": ContentType(parts=parts), + **gemini_params, + } + + return request_body + + +def process_embed_content_response( + input: EmbeddingInput, + model_response: EmbeddingResponse, + model: str, + response_json: dict, +) -> EmbeddingResponse: + """ + Process Gemini embedContent response (single embedding for multimodal input). + + Args: + input: Original input + model_response: EmbeddingResponse to populate + model: Model name + response_json: Raw JSON response from embedContent endpoint + + Returns: + EmbeddingResponse with single embedding + """ + if "embedding" not in response_json: + raise ValueError(f"embedContent response missing 'embedding' field: {response_json}") + + embedding_data = response_json["embedding"] + + openai_embedding = Embedding( + embedding=embedding_data["values"], + index=0, + object="embedding", + ) + + 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) + model_response.usage = Usage( + prompt_tokens=prompt_tokens, total_tokens=prompt_tokens + ) + + return model_response + + def process_response( input: EmbeddingInput, model_response: EmbeddingResponse, diff --git a/litellm/types/llms/vertex_ai.py b/litellm/types/llms/vertex_ai.py index 190e680b7b9..81de09595af 100644 --- a/litellm/types/llms/vertex_ai.py +++ b/litellm/types/llms/vertex_ai.py @@ -556,6 +556,17 @@ class VertexAIBatchEmbeddingsResponseObject(TypedDict): embeddings: List[ContentEmbeddings] +class GeminiEmbedContentRequestBody(TypedDict, total=False): + content: Required[ContentType] + taskType: TaskTypeEnum + title: str + outputDimensionality: int + + +class GeminiEmbedContentResponseObject(TypedDict): + embedding: ContentEmbeddings + + # Vertex AI Batch Prediction 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 7047be4241b..ba741d6c2bc 100644 --- a/tests/litellm/llms/vertex_ai/test_gemini_batch_embeddings.py +++ b/tests/litellm/llms/vertex_ai/test_gemini_batch_embeddings.py @@ -15,8 +15,16 @@ from unittest.mock import MagicMock, patch sys.path.insert(0, os.path.abspath("../../../..")) import pytest + import litellm from litellm.llms.custom_httpx.http_handler import HTTPHandler +from litellm.llms.vertex_ai.gemini_embeddings.batch_embed_content_transformation import ( + _is_multimodal_input, + _parse_data_url, + process_embed_content_response, + transform_openai_input_gemini_embed_content, +) +from litellm.types.utils import EmbeddingResponse def test_gemini_batch_embeddings_with_custom_api_base_and_auth_header(): @@ -47,11 +55,9 @@ def test_gemini_batch_embeddings_with_custom_api_base_and_auth_header(): mock_response = MagicMock() mock_response.status_code = 200 mock_response.json.return_value = { - "predictions": [ + "embeddings": [ { - "embeddings": { - "values": [0.1, 0.2, 0.3, 0.4, 0.5] - } + "values": [0.1, 0.2, 0.3, 0.4, 0.5] } ] } @@ -109,11 +115,9 @@ def test_gemini_batch_embeddings_with_extra_headers(): mock_response = MagicMock() mock_response.status_code = 200 mock_response.json.return_value = { - "predictions": [ + "embeddings": [ { - "embeddings": { - "values": [0.1, 0.2, 0.3] - } + "values": [0.1, 0.2, 0.3] } ] } @@ -143,3 +147,247 @@ def test_gemini_batch_embeddings_with_extra_headers(): assert "X-Custom" in headers assert headers["X-Custom"] == "custom-value" + +def test_is_multimodal_input_detection(): + """Test that _is_multimodal_input correctly detects multimodal inputs.""" + assert _is_multimodal_input("plain text") is False + assert _is_multimodal_input(["text1", "text2"]) is False + + assert _is_multimodal_input("data:image/png;base64,iVBORw0KGgo=") is True + assert _is_multimodal_input(["text", "data:image/png;base64,abc"]) is True + + assert _is_multimodal_input("files/abc123") is True + assert _is_multimodal_input(["text", "files/myfile"]) is True + + +def test_parse_data_url(): + """Test that _parse_data_url correctly extracts MIME type and base64 data.""" + mime_type, base64_data = _parse_data_url("data:image/png;base64,iVBORw0KGgo=") + assert mime_type == "image/png" + assert base64_data == "iVBORw0KGgo=" + + mime_type, base64_data = _parse_data_url("data:audio/mpeg;base64,SUQzBAA=") + assert mime_type == "audio/mpeg" + assert base64_data == "SUQzBAA=" + + mime_type, base64_data = _parse_data_url("data:video/mp4;base64,AAAAIGZ0eXA=") + assert mime_type == "video/mp4" + assert base64_data == "AAAAIGZ0eXA=" + + mime_type, base64_data = _parse_data_url("data:application/pdf;base64,JVBERi0=") + assert mime_type == "application/pdf" + assert base64_data == "JVBERi0=" + + +def test_mime_type_validation(): + """Test that unsupported MIME types raise ValueError.""" + with pytest.raises(ValueError, match="Unsupported MIME type"): + _parse_data_url("data:text/plain;base64,SGVsbG8=") + + with pytest.raises(ValueError, match="Unsupported MIME type"): + _parse_data_url("data:application/json;base64,e30=") + + +def test_parse_data_url_invalid_format(): + """Test that invalid data URL formats raise ValueError.""" + with pytest.raises(ValueError, match="Invalid data URL format"): + _parse_data_url("not-a-data-url") + + with pytest.raises(ValueError, match="missing comma"): + _parse_data_url("data:image/png;base64") + + +def test_transform_multimodal_text_and_image(): + """Test transformation of mixed text and image input.""" + input_data = [ + "The food was delicious", + "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" + ] + + result = transform_openai_input_gemini_embed_content( + input=input_data, + model="gemini-embedding-2-preview", + optional_params={}, + resolved_files=None, + ) + + assert "content" in result + assert "parts" in result["content"] + parts = result["content"]["parts"] + + assert len(parts) == 2 + assert parts[0]["text"] == "The food was delicious" + assert "inline_data" in parts[1] + assert parts[1]["inline_data"]["mime_type"] == "image/png" + assert "data" in parts[1]["inline_data"] + + +def test_transform_multimodal_with_file_reference(): + """Test transformation with Gemini file reference.""" + input_data = ["Some text", "files/abc123"] + + resolved_files = { + "files/abc123": { + "mime_type": "image/jpeg", + "uri": "https://generativelanguage.googleapis.com/v1beta/files/abc123" + } + } + + result = transform_openai_input_gemini_embed_content( + input=input_data, + model="gemini-embedding-2-preview", + optional_params={}, + resolved_files=resolved_files, + ) + + assert "content" in result + parts = result["content"]["parts"] + + assert len(parts) == 2 + assert parts[0]["text"] == "Some text" + assert "file_data" in parts[1] + assert parts[1]["file_data"]["mime_type"] == "image/jpeg" + assert parts[1]["file_data"]["file_uri"] == "https://generativelanguage.googleapis.com/v1beta/files/abc123" + + +def test_embed_content_response_processing(): + """Test processing of embedContent response (single embedding).""" + response_json = { + "embedding": { + "values": [0.1, 0.2, 0.3, 0.4, 0.5] + } + } + + model_response = EmbeddingResponse() + result = process_embed_content_response( + input=["test input"], + model_response=model_response, + model="gemini-embedding-2-preview", + response_json=response_json, + ) + + assert len(result.data) == 1 + assert result.data[0].embedding == [0.1, 0.2, 0.3, 0.4, 0.5] + assert result.data[0].index == 0 + assert result.data[0].object == "embedding" + assert result.model == "gemini-embedding-2-preview" + + +def test_gemini_multimodal_embedding_e2e(): + """Test end-to-end multimodal embedding call through litellm.embedding().""" + client = HTTPHandler() + + def mock_auth_token(*args, **kwargs): + return None, "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 = ( + {"x-goog-api-key": "test-key"}, + "https://generativelanguage.googleapis.com/v1beta/models/gemini-embedding-2-preview:embedContent?key=test-key" + ) + + 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="gemini/gemini-embedding-2-preview", + input=["The food was delicious", "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg=="], + api_key="test-key", + client=client + ) + + mock_post.assert_called_once() + + call_args = mock_post.call_args + kwargs = call_args.kwargs if hasattr(call_args, 'kwargs') else call_args[1] + + request_body = json.loads(kwargs.get("data", "{}")) + + assert "content" in request_body + assert "parts" in request_body["content"] + parts = request_body["content"]["parts"] + + assert len(parts) == 2 + assert parts[0]["text"] == "The food was delicious" + assert "inline_data" in parts[1] + assert parts[1]["inline_data"]["mime_type"] == "image/png" + + assert len(response.data) == 1 + assert response.data[0].embedding == [0.1, 0.2, 0.3, 0.4, 0.5] + + +def test_gemini_multimodal_embedding_with_audio(): + """Test multimodal embedding with audio input.""" + input_data = ["Audio description", "data:audio/mpeg;base64,SUQzBAAAAAA="] + + result = transform_openai_input_gemini_embed_content( + input=input_data, + model="gemini-embedding-2-preview", + optional_params={}, + resolved_files=None, + ) + + parts = result["content"]["parts"] + assert len(parts) == 2 + assert parts[0]["text"] == "Audio description" + assert parts[1]["inline_data"]["mime_type"] == "audio/mpeg" + + +def test_gemini_multimodal_embedding_with_video(): + """Test multimodal embedding with video input.""" + input_data = ["data:video/mp4;base64,AAAAIGZ0eXBpc29tAAACAGlzb21pc28yYXZjMW1wNDEAAAAIZnJlZQAA"] + + result = transform_openai_input_gemini_embed_content( + input=input_data, + model="gemini-embedding-2-preview", + optional_params={}, + resolved_files=None, + ) + + parts = result["content"]["parts"] + assert len(parts) == 1 + assert parts[0]["inline_data"]["mime_type"] == "video/mp4" + + + +def test_transform_with_optional_params(): + """Test that optional params like outputDimensionality are passed through.""" + input_data = ["test text"] + + result = transform_openai_input_gemini_embed_content( + input=input_data, + model="gemini-embedding-2-preview", + optional_params={"outputDimensionality": 768, "taskType": "SEMANTIC_SIMILARITY"}, + resolved_files=None, + ) + + assert result["outputDimensionality"] == 768 + assert result["taskType"] == "SEMANTIC_SIMILARITY" + + +def test_dimensions_mapped_to_output_dimensionality(): + """Test that OpenAI 'dimensions' param is mapped to Gemini 'outputDimensionality'.""" + input_data = ["test text"] + + result = transform_openai_input_gemini_embed_content( + input=input_data, + model="gemini-embedding-2-preview", + optional_params={"dimensions": 768}, + resolved_files=None, + ) + + assert "outputDimensionality" in result + assert result["outputDimensionality"] == 768 + assert "dimensions" not in result + From 2c4a495619fd95cf203bb9b60485ebf6fb2a2377 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Wed, 11 Mar 2026 10:39:34 +0530 Subject: [PATCH 2/7] 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 3c5cbb65437..c02d63414c5 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 364519e1fe3..529b998810f 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 b5d5c06924d..52221c47deb 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 4367ec789b3..eebcdb3ec59 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}") From d25b8e6d009a52dce7d6ed0906ccee8deedc905d Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Wed, 11 Mar 2026 10:58:04 +0530 Subject: [PATCH 3/7] Add support for gcs url for vertex ai embeddings --- .../batch_embed_content_transformation.py | 57 ++++++++++++++-- .../vertex_ai/test_gemini_batch_embeddings.py | 65 +++++++++++++++++++ 2 files changed, 118 insertions(+), 4 deletions(-) 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 6070c70677b..b2bf2c6eb5c 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 @@ -35,6 +35,46 @@ def _is_file_reference(s: str) -> bool: return isinstance(s, str) and s.startswith("files/") +def _is_gcs_url(s: str) -> bool: + """Check if string is a GCS URL (gs://...).""" + return isinstance(s, str) and s.startswith("gs://") + + +def _infer_mime_type_from_gcs_url(gcs_url: str) -> str: + """ + Infer MIME type from GCS URL file extension. + + Args: + gcs_url: GCS URL like gs://bucket/path/to/file.png + + Returns: + str: Inferred MIME type + + Raises: + ValueError: If file extension is not supported + """ + extension_to_mime = { + ".png": "image/png", + ".jpg": "image/jpeg", + ".jpeg": "image/jpeg", + ".mp3": "audio/mpeg", + ".wav": "audio/wav", + ".mp4": "video/mp4", + ".mov": "video/quicktime", + ".pdf": "application/pdf", + } + + gcs_url_lower = gcs_url.lower() + for ext, mime_type in extension_to_mime.items(): + if gcs_url_lower.endswith(ext): + return mime_type + + raise ValueError( + f"Unable to infer MIME type from GCS URL: {gcs_url}. " + f"Supported extensions: {', '.join(extension_to_mime.keys())}" + ) + + def _parse_data_url(data_url: str) -> Tuple[str, str]: """ Parse a data URL to extract the media type and base64 data. @@ -76,13 +116,13 @@ def _parse_data_url(data_url: str) -> Tuple[str, str]: def _is_multimodal_input(input: EmbeddingInput) -> bool: """ - Check if the input contains multimodal data (data URIs or file references). + Check if the input contains multimodal data (data URIs, file references, or GCS URLs). Args: input: EmbeddingInput (str or List[str]) Returns: - bool: True if any element is a data URI or file reference + bool: True if any element is a data URI, file reference, or GCS URL """ if isinstance(input, str): input_list = [input] @@ -95,6 +135,8 @@ def _is_multimodal_input(input: EmbeddingInput) -> bool: return True if _is_file_reference(element): return True + if _is_gcs_url(element): + return True return False @@ -166,15 +208,22 @@ def transform_openai_input_gemini_embed_content( mime_type, base64_data = _parse_data_url(element) blob: BlobType = {"mime_type": mime_type, "data": base64_data} parts.append(PartType(inline_data=blob)) + elif _is_gcs_url(element): + mime_type = _infer_mime_type_from_gcs_url(element) + file_data: FileDataType = { + "mime_type": mime_type, + "file_uri": element, + } + parts.append(PartType(file_data=file_data)) elif _is_file_reference(element): if element not in resolved_files: raise ValueError(f"File reference {element} not resolved") file_info = resolved_files[element] - file_data: FileDataType = { + file_data_ref: FileDataType = { "mime_type": file_info["mime_type"], "file_uri": file_info["uri"], } - parts.append(PartType(file_data=file_data)) + parts.append(PartType(file_data=file_data_ref)) else: parts.append(PartType(text=element)) 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 ba741d6c2bc..302facefcb7 100644 --- a/tests/litellm/llms/vertex_ai/test_gemini_batch_embeddings.py +++ b/tests/litellm/llms/vertex_ai/test_gemini_batch_embeddings.py @@ -391,3 +391,68 @@ def test_dimensions_mapped_to_output_dimensionality(): assert result["outputDimensionality"] == 768 assert "dimensions" not in result + +def test_is_gcs_url(): + """Test GCS URL detection.""" + from litellm.llms.vertex_ai.gemini_embeddings.batch_embed_content_transformation import ( + _is_gcs_url, + ) + + assert _is_gcs_url("gs://my-bucket/path/to/file.png") is True + assert _is_gcs_url("gs://bucket/image.jpg") is True + assert _is_gcs_url("https://storage.googleapis.com/bucket/file.png") is False + assert _is_gcs_url("files/abc123") is False + assert _is_gcs_url("data:image/png;base64,abc") is False + assert _is_gcs_url("regular text") is False + + +def test_infer_mime_type_from_gcs_url(): + """Test MIME type inference from GCS URL.""" + from litellm.llms.vertex_ai.gemini_embeddings.batch_embed_content_transformation import ( + _infer_mime_type_from_gcs_url, + ) + + assert _infer_mime_type_from_gcs_url("gs://bucket/image.png") == "image/png" + assert _infer_mime_type_from_gcs_url("gs://bucket/photo.jpg") == "image/jpeg" + assert _infer_mime_type_from_gcs_url("gs://bucket/photo.JPEG") == "image/jpeg" + assert _infer_mime_type_from_gcs_url("gs://bucket/audio.mp3") == "audio/mpeg" + assert _infer_mime_type_from_gcs_url("gs://bucket/audio.wav") == "audio/wav" + assert _infer_mime_type_from_gcs_url("gs://bucket/video.mp4") == "video/mp4" + assert _infer_mime_type_from_gcs_url("gs://bucket/video.mov") == "video/quicktime" + assert _infer_mime_type_from_gcs_url("gs://bucket/doc.pdf") == "application/pdf" + + with pytest.raises(ValueError, match="Unable to infer MIME type"): + _infer_mime_type_from_gcs_url("gs://bucket/file.txt") + + +def test_transform_multimodal_with_gcs_url(): + """Test transformation with GCS URL.""" + input_data = [ + "Describe this image", + "gs://my-bucket/images/photo.png" + ] + + result = transform_openai_input_gemini_embed_content( + input=input_data, + model="gemini-embedding-2-preview", + optional_params={}, + resolved_files=None, + ) + + parts = result["content"]["parts"] + assert len(parts) == 2 + assert parts[0]["text"] == "Describe this image" + assert parts[1]["file_data"]["mime_type"] == "image/png" + assert parts[1]["file_data"]["file_uri"] == "gs://my-bucket/images/photo.png" + + +def test_multimodal_input_detection_with_gcs(): + """Test that GCS URLs are detected as multimodal.""" + from litellm.llms.vertex_ai.gemini_embeddings.batch_embed_content_transformation import ( + _is_multimodal_input, + ) + + assert _is_multimodal_input(["text", "gs://bucket/file.png"]) is True + assert _is_multimodal_input("gs://bucket/video.mp4") is True + assert _is_multimodal_input(["just text", "more text"]) is False + From 8c5478df705b05c52aa5c2958b75822e3bf3f9cf Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Wed, 11 Mar 2026 10:58:50 +0530 Subject: [PATCH 4/7] Add embedding model in model map --- ...odel_prices_and_context_window_backup.json | 39 +++++++++++++++++++ model_prices_and_context_window.json | 39 +++++++++++++++++++ 2 files changed, 78 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index d3dd6b3d992..b53e1e14d7d 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -15962,6 +15962,32 @@ "output_vector_size": 3072, "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models" }, + "gemini-embedding-2-preview": { + "input_cost_per_audio_per_second": 0.00016, + "input_cost_per_image": 0.00012, + "input_cost_per_token": 2e-07, + "input_cost_per_video_per_second": 0.0237, + "litellm_provider": "vertex_ai-embedding-models", + "max_input_tokens": 8192, + "max_tokens": 8192, + "mode": "embedding", + "output_cost_per_token": 0, + "output_vector_size": 3072, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", + "uses_embed_content": true + }, + "vertex_ai/gemini-embedding-2-preview": { + "input_cost_per_token": 1.5e-07, + "litellm_provider": "vertex_ai", + "max_input_tokens": 8192, + "max_tokens": 8192, + "mode": "embedding", + "output_cost_per_token": 0, + "output_vector_size": 3072, + "source": "https://ai.google.dev/gemini-api/docs/embeddings#multimodal", + "supports_multimodal": true, + "uses_embed_content": true + }, "gemini-flash-experimental": { "input_cost_per_character": 0, "input_cost_per_token": 0, @@ -16039,6 +16065,19 @@ "source": "https://ai.google.dev/gemini-api/docs/embeddings#model-versions", "tpm": 10000000 }, + "gemini/gemini-embedding-2-preview": { + "input_cost_per_token": 1.5e-07, + "litellm_provider": "gemini", + "max_input_tokens": 8192, + "max_tokens": 8192, + "mode": "embedding", + "output_cost_per_token": 0, + "output_vector_size": 3072, + "rpm": 10000, + "source": "https://ai.google.dev/gemini-api/docs/embeddings#multimodal", + "supports_multimodal": true, + "tpm": 10000000 + }, "gemini/gemini-1.5-flash": { "deprecation_date": "2025-09-29", "input_cost_per_token": 7.5e-08, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index d3dd6b3d992..b53e1e14d7d 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -15962,6 +15962,32 @@ "output_vector_size": 3072, "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models" }, + "gemini-embedding-2-preview": { + "input_cost_per_audio_per_second": 0.00016, + "input_cost_per_image": 0.00012, + "input_cost_per_token": 2e-07, + "input_cost_per_video_per_second": 0.0237, + "litellm_provider": "vertex_ai-embedding-models", + "max_input_tokens": 8192, + "max_tokens": 8192, + "mode": "embedding", + "output_cost_per_token": 0, + "output_vector_size": 3072, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", + "uses_embed_content": true + }, + "vertex_ai/gemini-embedding-2-preview": { + "input_cost_per_token": 1.5e-07, + "litellm_provider": "vertex_ai", + "max_input_tokens": 8192, + "max_tokens": 8192, + "mode": "embedding", + "output_cost_per_token": 0, + "output_vector_size": 3072, + "source": "https://ai.google.dev/gemini-api/docs/embeddings#multimodal", + "supports_multimodal": true, + "uses_embed_content": true + }, "gemini-flash-experimental": { "input_cost_per_character": 0, "input_cost_per_token": 0, @@ -16039,6 +16065,19 @@ "source": "https://ai.google.dev/gemini-api/docs/embeddings#model-versions", "tpm": 10000000 }, + "gemini/gemini-embedding-2-preview": { + "input_cost_per_token": 1.5e-07, + "litellm_provider": "gemini", + "max_input_tokens": 8192, + "max_tokens": 8192, + "mode": "embedding", + "output_cost_per_token": 0, + "output_vector_size": 3072, + "rpm": 10000, + "source": "https://ai.google.dev/gemini-api/docs/embeddings#multimodal", + "supports_multimodal": true, + "tpm": 10000000 + }, "gemini/gemini-1.5-flash": { "deprecation_date": "2025-09-29", "input_cost_per_token": 7.5e-08, From 1c144fc8961a850d8e84ca7c2feff1906bb9fab3 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Wed, 11 Mar 2026 11:02:49 +0530 Subject: [PATCH 5/7] Add embedding model documentation --- .../gemini_embedding_2_multimodal/index.md | 169 ++++++++++++++++++ .../docs/embedding/supported_embedding.md | 51 ++++++ .../docs/providers/vertex_embedding.md | 66 +++++++ 3 files changed, 286 insertions(+) create mode 100644 docs/my-website/blog/gemini_embedding_2_multimodal/index.md diff --git a/docs/my-website/blog/gemini_embedding_2_multimodal/index.md b/docs/my-website/blog/gemini_embedding_2_multimodal/index.md new file mode 100644 index 00000000000..8c09432e3b6 --- /dev/null +++ b/docs/my-website/blog/gemini_embedding_2_multimodal/index.md @@ -0,0 +1,169 @@ +--- +slug: gemini_embedding_2_multimodal +title: "Gemini Embedding 2 Preview: Multimodal Embeddings on LiteLLM" +date: 2025-03-11T10:00:00 +authors: + - name: Sameer Kankute + title: SWE @ LiteLLM (LLM Translation) + url: https://www.linkedin.com/in/sameer-kankute/ + image_url: https://pbs.twimg.com/profile_images/2001352686994907136/ONgNuSk5_400x400.jpg +description: "Generate embeddings from text, images, audio, video, and PDFs with gemini-embedding-2-preview on LiteLLM via Gemini API and Vertex AI." +tags: [gemini, embeddings, multimodal, vertex ai] +hide_table_of_contents: false +--- + +import Tabs from '@theme/Tabs'; +import TabItem from '@theme/TabItem'; + +# Gemini Embedding 2 Preview: Multimodal Embeddings + +LiteLLM now supports **multimodal embeddings** with `gemini-embedding-2-preview`—generating a single embedding from a mix of text, images, audio, video, and PDF content. Available via both the **Gemini API** (API key) and **Vertex AI** (GCP credentials). + +## Supported Input Types + +| Modality | Supported Formats | +|----------|-------------------| +| **Text** | Plain text | +| **Image** | PNG, JPEG | +| **Audio** | MP3, WAV | +| **Video** | MP4, MOV | +| **Documents** | PDF | + +## Input Formats + +LiteLLM accepts three input formats for multimodal content: + +1. **Data URIs** – Base64-encoded inline: `data:image/png;base64,` +2. **GCS URLs** – Cloud Storage paths (Vertex AI): `gs://bucket/path/to/file.png` +3. **Gemini File References** – Pre-uploaded files (Gemini API): `files/abc123` + +## Quick Start + + + + +```python +from litellm import embedding +import os + +os.environ["GEMINI_API_KEY"] = "your-api-key" + +# Text + Image (base64) +response = embedding( + model="gemini/gemini-embedding-2-preview", + input=[ + "The food was delicious and the waiter...", + "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAgAAAAIAQMAAAD+wSzIAAAABlBMVEX///+/v7+jQ3Y5AAAADklEQVQI12P4AIX8EAgALgAD/aNpbtEAAAAASUVORK5CYII" + ], +) +print(response) +``` + + + + + +```python +import litellm +from litellm import embedding + +litellm.vertex_project = "your-project-id" +litellm.vertex_location = "us-central1" + +# Text + Image (GCS URL) +response = embedding( + model="vertex_ai/gemini-embedding-2-preview", + input=[ + "Describe this image", + "gs://my-bucket/images/photo.png" + ], +) +print(response) +``` + + + + + +**1. Config (config.yaml)** + +```yaml +model_list: + - model_name: gemini-embedding-2-preview + litellm_params: + model: gemini/gemini-embedding-2-preview + api_key: os.environ/GEMINI_API_KEY + - model_name: vertex-gemini-embedding-2-preview + litellm_params: + model: vertex_ai/gemini-embedding-2-preview + vertex_project: os.environ/VERTEXAI_PROJECT + vertex_location: os.environ/VERTEXAI_LOCATION + +general_settings: + master_key: sk-1234 +``` + +**2. Start proxy** + +```bash +litellm --config config.yaml +``` + +**3. Call embeddings** + +```bash +curl -X POST http://localhost:4000/embeddings \ + -H "Authorization: Bearer sk-1234" \ + -H "Content-Type: application/json" \ + -d '{ + "model": "gemini-embedding-2-preview", + "input": [ + "The food was delicious and the waiter...", + "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAgAAAAIAQMAAAD+wSzIAAAABlBMVEX///+/v7+jQ3Y5AAAADklEQVQI12P4AIX8EAgALgAD/aNpbtEAAAAASUVORK5CYII" + ] + }' +``` + + + + +## Input Format Examples + +| Format | Example | Provider | +|--------|---------|----------| +| **Data URI** | `data:image/png;base64,...` | Gemini, Vertex AI | +| **GCS URL** | `gs://bucket/path/image.png` | Vertex AI | +| **File reference** | `files/abc123` | Gemini API only | + +### Supported MIME Types for Data URIs + +- **Images:** `image/png`, `image/jpeg` +- **Audio:** `audio/mpeg`, `audio/wav` +- **Video:** `video/mp4`, `video/quicktime` +- **Documents:** `application/pdf` + +### GCS URL MIME Inference + +For Vertex AI, MIME types are inferred from file extensions: + +- `.png` → `image/png` +- `.jpg` / `.jpeg` → `image/jpeg` +- `.mp3` → `audio/mpeg` +- `.wav` → `audio/wav` +- `.mp4` → `video/mp4` +- `.mov` → `video/quicktime` +- `.pdf` → `application/pdf` + +## Optional Parameters + +| Parameter | Description | Maps to | +|-----------|-------------|---------| +| `dimensions` | Output embedding size | `outputDimensionality` | + +```python +response = embedding( + model="gemini/gemini-embedding-2-preview", + input=["text to embed"], + dimensions=768, # Optional: control output vector size +) +``` diff --git a/docs/my-website/docs/embedding/supported_embedding.md b/docs/my-website/docs/embedding/supported_embedding.md index 11ca4da48a4..87acd0b33a5 100644 --- a/docs/my-website/docs/embedding/supported_embedding.md +++ b/docs/my-website/docs/embedding/supported_embedding.md @@ -514,6 +514,57 @@ All models listed [here](https://ai.google.dev/gemini-api/docs/models/gemini) ar | Model Name | Function Call | | :--- | :--- | | text-embedding-004 | `embedding(model="gemini/text-embedding-004", input)` | +| gemini-embedding-2-preview | `embedding(model="gemini/gemini-embedding-2-preview", input)` | [Multimodal docs](#gemini-embedding-2-preview-multimodal) | + +### Gemini Embedding 2 Preview (Multimodal) + +`gemini-embedding-2-preview` supports **multimodal embeddings**—text, images, audio, video, and PDF in a single request. See [blog post](/blog/gemini_embedding_2_multimodal) for details. + +**Input formats:** +- **Data URIs:** `data:image/png;base64,` +- **Gemini file references:** `files/abc123` (pre-uploaded via Gemini Files API) + +**Supported MIME types:** `image/png`, `image/jpeg`, `audio/mpeg`, `audio/wav`, `video/mp4`, `video/quicktime`, `application/pdf` + + + + +```python +from litellm import embedding +import os +os.environ["GEMINI_API_KEY"] = "" + +# Text + Image (base64) +response = embedding( + model="gemini/gemini-embedding-2-preview", + input=[ + "The food was delicious and the waiter...", + "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAgAAAAIAQMAAAD+wSzIAAAABlBMVEX///+/v7+jQ3Y5AAAADklEQVQI12P4AIX8EAgALgAD/aNpbtEAAAAASUVORK5CYII" + ], +) +print(response) +``` + + + + +```bash +curl -X POST http://localhost:4000/embeddings \ + -H "Authorization: Bearer sk-1234" \ + -H "Content-Type: application/json" \ + -d '{ + "model": "gemini-embedding-2-preview", + "input": [ + "The food was delicious and the waiter...", + "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAgAAAAIAQMAAAD+wSzIAAAABlBMVEX///+/v7+jQ3Y5AAAADklEQVQI12P4AIX8EAgALgAD/aNpbtEAAAAASUVORK5CYII" + ] + }' +``` + + + + +**Optional:** `dimensions` maps to Gemini's `outputDimensionality`. ## Vertex AI Embedding Models diff --git a/docs/my-website/docs/providers/vertex_embedding.md b/docs/my-website/docs/providers/vertex_embedding.md index 5656ade337b..9b530f2ae06 100644 --- a/docs/my-website/docs/providers/vertex_embedding.md +++ b/docs/my-website/docs/providers/vertex_embedding.md @@ -79,6 +79,7 @@ All models listed [here](https://github.com/BerriAI/litellm/blob/57f37f743886a02 | textembedding-gecko@003 | `embedding(model="vertex_ai/textembedding-gecko@003", input)` | | text-embedding-preview-0409 | `embedding(model="vertex_ai/text-embedding-preview-0409", input)` | | text-multilingual-embedding-preview-0409 | `embedding(model="vertex_ai/text-multilingual-embedding-preview-0409", input)` | +| gemini-embedding-2-preview | `embedding(model="vertex_ai/gemini-embedding-2-preview", input)` | [Multimodal docs](#gemini-embedding-2-preview-multimodal) | | Fine-tuned OR Custom Embedding models | `embedding(model="vertex_ai/", input)` | ### Supported OpenAI (Unified) Params @@ -257,6 +258,71 @@ model_list: ## **Multi-Modal Embeddings** +### Gemini Embedding 2 Preview (Multimodal) + +`gemini-embedding-2-preview` supports **unified multimodal embeddings**—text, images, audio, video, and PDF in a single request. See [blog post](/blog/gemini_embedding_2_multimodal) for details. + +**Input formats:** +- **Data URIs:** `data:image/png;base64,` +- **GCS URLs:** `gs://bucket/path/to/file.png` (MIME type inferred from extension) + +**Supported MIME types:** `image/png`, `image/jpeg`, `audio/mpeg`, `audio/wav`, `video/mp4`, `video/quicktime`, `application/pdf` + + + + +```python +import litellm +from litellm import embedding + +litellm.vertex_project = "your-project-id" +litellm.vertex_location = "us-central1" + +# Text + Image (GCS URL) +response = embedding( + model="vertex_ai/gemini-embedding-2-preview", + input=[ + "Describe this image", + "gs://my-bucket/images/photo.png" + ], +) + +# Text + Image (base64) +response = embedding( + model="vertex_ai/gemini-embedding-2-preview", + input=[ + "The food was delicious", + "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAgAAAAIAQMAAAD+wSzIAAAABlBMVEX///+/v7+jQ3Y5AAAADklEQVQI12P4AIX8EAgALgAD/aNpbtEAAAAASUVORK5CYII" + ], +) +``` + + + + +```yaml +model_list: + - model_name: vertex-gemini-embedding-2-preview + litellm_params: + model: vertex_ai/gemini-embedding-2-preview + vertex_project: "your-project-id" + vertex_location: "us-central1" +``` + +```bash +curl -X POST http://localhost:4000/embeddings \ + -H "Authorization: Bearer sk-1234" \ + -H "Content-Type: application/json" \ + -d '{ + "model": "vertex-gemini-embedding-2-preview", + "input": ["Describe this", "gs://bucket/image.png"] + }' +``` + + + + +### multimodalembedding@001 (Legacy) Known Limitations: - Only supports 1 image / video / image per request From 2a9bcf2530e61ef573db9d6955586caa06789eed Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Wed, 11 Mar 2026 11:41:29 +0530 Subject: [PATCH 6/7] 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 1447eb4b92a..25c3465807c 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 b2bf2c6eb5c..41f477d9db9 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 302facefcb7..1ed1de01b5f 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 + From e394914d34d962c8b348b09652bf9d866a99090b Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Wed, 11 Mar 2026 11:49:02 +0530 Subject: [PATCH 7/7] Fix code qa --- .../vertex_ai/gemini_embeddings/batch_embed_content_handler.py | 3 +-- 1 file changed, 1 insertion(+), 2 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 25c3465807c..68901340c7c 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 @@ -3,7 +3,7 @@ Google AI Studio /batchEmbedContents Embeddings Endpoint """ import json -from typing import Any, Dict, List, Literal, Optional, Union +from typing import Any, Dict, Literal, Optional, Union import httpx @@ -15,7 +15,6 @@ from litellm.llms.custom_httpx.http_handler import ( ) from litellm.types.llms.openai import EmbeddingInput from litellm.types.llms.vertex_ai import ( - GeminiEmbedContentResponseObject, VertexAIBatchEmbeddingsRequestBody, VertexAIBatchEmbeddingsResponseObject, )