mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-09 22:31:41 +00:00
Add support for gemini multimodal embedings
This commit is contained in:
Sameer Kankute
parent
40210ce750
commit
b108c02fd7
4 changed files with 656 additions and 63 deletions
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue