Merge pull request #24370 from Chesars/fix/gemini-embedding-drop-unsupported-params

fix(gemini): filter unsupported params from embedding requests

litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py

@@ -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] = []

tests/litellm/llms/vertex_ai/test_gemini_batch_embeddings.py

@@ -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
+