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 +