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 1743122753..99165c37c9 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 @@ -13,8 +13,8 @@ from litellm.llms.custom_httpx.http_handler import ( HTTPHandler, get_async_httpx_client, ) -from litellm.types.llms.openai import EmbeddingInput from litellm.types.llms.vertex_ai import ( + GeminiEmbeddingInput, VertexAIBatchEmbeddingsRequestBody, VertexAIBatchEmbeddingsResponseObject, ) @@ -33,7 +33,7 @@ from .batch_embed_content_transformation import ( class GoogleBatchEmbeddings(VertexLLM): @staticmethod def _flatten_and_detect_file_refs( - input: EmbeddingInput, + input: GeminiEmbeddingInput, ) -> Tuple[List[str], bool]: """Flatten nested input lists and detect file references.""" input_list = [input] if isinstance(input, str) else input @@ -48,7 +48,7 @@ class GoogleBatchEmbeddings(VertexLLM): def _resolve_file_references( self, - input: Union[EmbeddingInput, List[str]], + input: GeminiEmbeddingInput, api_key: str, sync_handler: HTTPHandler, ) -> Dict[str, Dict[str, str]]: @@ -56,7 +56,7 @@ class GoogleBatchEmbeddings(VertexLLM): Resolve Gemini file references (files/...) to get mime_type and uri. Args: - input: EmbeddingInput that may contain file references + input: GeminiEmbeddingInput that may contain file references api_key: Gemini API key sync_handler: HTTP client @@ -87,7 +87,7 @@ class GoogleBatchEmbeddings(VertexLLM): async def _async_resolve_file_references( self, - input: Union[EmbeddingInput, List[str]], + input: GeminiEmbeddingInput, api_key: str, async_handler: AsyncHTTPHandler, ) -> Dict[str, Dict[str, str]]: @@ -95,7 +95,7 @@ class GoogleBatchEmbeddings(VertexLLM): Async version of _resolve_file_references. Args: - input: EmbeddingInput that may contain file references + input: GeminiEmbeddingInput that may contain file references api_key: Gemini API key async_handler: Async HTTP client @@ -127,7 +127,7 @@ class GoogleBatchEmbeddings(VertexLLM): def batch_embeddings( # noqa: PLR0915 self, model: str, - input: EmbeddingInput, + input: GeminiEmbeddingInput, print_verbose, model_response: EmbeddingResponse, custom_llm_provider: Literal["gemini", "vertex_ai"], @@ -291,7 +291,7 @@ class GoogleBatchEmbeddings(VertexLLM): url: str, data: Optional[Union[VertexAIBatchEmbeddingsRequestBody, dict]], model_response: EmbeddingResponse, - input: EmbeddingInput, + input: GeminiEmbeddingInput, timeout: Optional[Union[float, httpx.Timeout]], headers={}, client: Optional[AsyncHTTPHandler] = None, 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 11457fc924..e1b365c9f4 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 @@ -6,12 +6,12 @@ Why separate file? Make it easy to see how transformation works from typing import Dict, List, Optional, Tuple -from litellm.types.llms.openai import EmbeddingInput from litellm.types.llms.vertex_ai import ( BlobType, ContentType, EmbedContentRequest, FileDataType, + GeminiEmbeddingInput, PartType, VertexAIBatchEmbeddingsRequestBody, VertexAIBatchEmbeddingsResponseObject, @@ -114,13 +114,13 @@ def _parse_data_url(data_url: str) -> Tuple[str, str]: return media_type, base64_data -def _is_multimodal_input(input: EmbeddingInput) -> bool: +def _is_multimodal_input(input: GeminiEmbeddingInput) -> bool: """ Check if the input contains multimodal data (data URIs, file references, GCS URLs, or nested lists for combined embeddings). Args: - input: EmbeddingInput — str, List[str], or List[Union[str, List[str]]] + input: GeminiEmbeddingInput — str, List[str], or List[List[str]] for combined embeddings Returns: bool: True if any element is multimodal or a nested list @@ -199,7 +199,7 @@ def _filter_embed_params(optional_params: dict) -> dict: def transform_openai_input_gemini_content( - input: EmbeddingInput, + input: GeminiEmbeddingInput, model: str, optional_params: dict, resolved_files: Optional[Dict[str, Dict[str, str]]] = None, @@ -252,7 +252,7 @@ def transform_openai_input_gemini_content( def transform_openai_input_gemini_embed_content( - input: EmbeddingInput, + input: GeminiEmbeddingInput, model: str, optional_params: dict, resolved_files: Optional[Dict[str, Dict[str, str]]] = None, @@ -261,7 +261,7 @@ def transform_openai_input_gemini_embed_content( Transform OpenAI embedding input to Gemini embedContent format (multimodal). Args: - input: EmbeddingInput (str or List[str]) with text, data URIs, or file references + input: GeminiEmbeddingInput 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} @@ -295,7 +295,7 @@ def transform_openai_input_gemini_embed_content( def process_embed_content_response( - input: EmbeddingInput, + input: GeminiEmbeddingInput, model_response: EmbeddingResponse, model: str, response_json: dict, @@ -341,7 +341,7 @@ def process_embed_content_response( def process_response( - input: EmbeddingInput, + input: GeminiEmbeddingInput, model_response: EmbeddingResponse, model: str, _predictions: VertexAIBatchEmbeddingsResponseObject, diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index 12fd2284ed..2fd0c4ea97 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -103,7 +103,7 @@ FileTypes = Union[ ] -EmbeddingInput = Union[str, List[Union[str, List[str]]]] +EmbeddingInput = Union[str, List[str]] class HttpxBinaryResponseContent(_HttpxBinaryResponseContent): diff --git a/litellm/types/llms/vertex_ai.py b/litellm/types/llms/vertex_ai.py index 2e7d57cef2..87bf11a902 100644 --- a/litellm/types/llms/vertex_ai.py +++ b/litellm/types/llms/vertex_ai.py @@ -6,6 +6,13 @@ from typing_extensions import ( TypedDict, ) +from litellm.types.llms.openai import EmbeddingInput + +# Gemini supports nested-list inputs (e.g. [["text", "image"]]) as an explicit +# opt-in for combined embeddings — a provider-specific extension of the +# OpenAI-faithful EmbeddingInput shape. +GeminiEmbeddingInput = Union[EmbeddingInput, List[List[str]]] + class FunctionResponse(TypedDict): name: str