mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-17 23:51:30 +00:00
fix(gemini): filter unsupported params from embedding requests
The Gemini batch embedding transformation was spreading all optional_params into the request body via **gemini_params. Params like max_tokens (injected by add_provider_specific_params_to_optional_params) would reach the Gemini API and cause a 400 BadRequestError. Extract _filter_embed_params() that maps dimensions/task_type and keeps only the fields Gemini embeddings actually accept (outputDimensionality, taskType, title). Applied to both transform_openai_input_gemini_content and transform_openai_input_gemini_embed_content. This also fixes drop_params: true not preventing the error, since the param was re-injected after the drop_params check. Fixes #24293
This commit is contained in:
parent
f5194b5ce3
commit
db0d85eefd
2 changed files with 63 additions and 10 deletions
|
|
@ -141,6 +141,19 @@ def _is_multimodal_input(input: EmbeddingInput) -> bool:
|
|||
return False
|
||||
|
||||
|
||||
_SUPPORTED_EMBED_PARAMS = {"outputDimensionality", "taskType", "title"}
|
||||
|
||||
|
||||
def _filter_embed_params(optional_params: dict) -> dict:
|
||||
"""Map and filter optional_params to only include Gemini embedding fields."""
|
||||
gemini_params = optional_params.copy()
|
||||
if "dimensions" in gemini_params:
|
||||
gemini_params["outputDimensionality"] = gemini_params.pop("dimensions")
|
||||
if "task_type" in gemini_params:
|
||||
gemini_params["taskType"] = gemini_params.pop("task_type")
|
||||
return {k: v for k, v in gemini_params.items() if k in _SUPPORTED_EMBED_PARAMS}
|
||||
|
||||
|
||||
def transform_openai_input_gemini_content(
|
||||
input: EmbeddingInput, model: str, optional_params: dict
|
||||
) -> VertexAIBatchEmbeddingsRequestBody:
|
||||
|
|
@ -149,11 +162,7 @@ def transform_openai_input_gemini_content(
|
|||
"""
|
||||
gemini_model_name = "models/{}".format(model)
|
||||
|
||||
gemini_params = optional_params.copy()
|
||||
if "dimensions" in gemini_params:
|
||||
gemini_params["outputDimensionality"] = gemini_params.pop("dimensions")
|
||||
if "task_type" in gemini_params:
|
||||
gemini_params["taskType"] = gemini_params.pop("task_type")
|
||||
gemini_params = _filter_embed_params(optional_params)
|
||||
|
||||
requests: List[EmbedContentRequest] = []
|
||||
if isinstance(input, str):
|
||||
|
|
@ -195,11 +204,7 @@ def transform_openai_input_gemini_embed_content(
|
|||
"""
|
||||
resolved_files = resolved_files or {}
|
||||
|
||||
gemini_params = optional_params.copy()
|
||||
if "dimensions" in gemini_params:
|
||||
gemini_params["outputDimensionality"] = gemini_params.pop("dimensions")
|
||||
if "task_type" in gemini_params:
|
||||
gemini_params["taskType"] = gemini_params.pop("task_type")
|
||||
gemini_params = _filter_embed_params(optional_params)
|
||||
|
||||
input_list = [input] if isinstance(input, str) else input
|
||||
parts: List[PartType] = []
|
||||
|
|
|
|||
|
|
@ -19,6 +19,7 @@ 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 (
|
||||
_filter_embed_params,
|
||||
_is_multimodal_input,
|
||||
_parse_data_url,
|
||||
process_embed_content_response,
|
||||
|
|
@ -563,3 +564,50 @@ def test_vertex_ai_text_only_embedding_uses_embed_content():
|
|||
assert data["content"]["parts"][0]["text"] == "Hello, world!"
|
||||
assert len(response.data) == 1
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Unsupported params filtering tests (#24293)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_filter_embed_params_drops_unsupported():
|
||||
"""Unsupported params like max_tokens should be filtered out."""
|
||||
result = _filter_embed_params({"dimensions": 768, "max_tokens": 256, "temperature": 0.5})
|
||||
assert result == {"outputDimensionality": 768}
|
||||
|
||||
|
||||
def test_filter_embed_params_keeps_supported():
|
||||
"""All supported Gemini embedding params should pass through."""
|
||||
result = _filter_embed_params({
|
||||
"dimensions": 768,
|
||||
"task_type": "RETRIEVAL_DOCUMENT",
|
||||
"title": "My doc",
|
||||
})
|
||||
assert result == {
|
||||
"outputDimensionality": 768,
|
||||
"taskType": "RETRIEVAL_DOCUMENT",
|
||||
"title": "My doc",
|
||||
}
|
||||
|
||||
|
||||
def test_batch_embed_content_drops_max_tokens():
|
||||
"""max_tokens in optional_params should not appear in the batch request."""
|
||||
result = transform_openai_input_gemini_content(
|
||||
input="test text",
|
||||
model="text-embedding-004",
|
||||
optional_params={"max_tokens": 256},
|
||||
)
|
||||
for request in result["requests"]:
|
||||
assert "max_tokens" not in request
|
||||
|
||||
|
||||
def test_embed_content_drops_max_tokens():
|
||||
"""max_tokens in optional_params should not appear in the embedContent request."""
|
||||
result = transform_openai_input_gemini_embed_content(
|
||||
input=["test text"],
|
||||
model="gemini-embedding-001",
|
||||
optional_params={"max_tokens": 256},
|
||||
resolved_files=None,
|
||||
)
|
||||
assert "max_tokens" not in result
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue