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 08831a8215..c1733c7242 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 @@ -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] = [] 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 a8e427d3bc..5f4c7b564c 100644 --- a/tests/litellm/llms/vertex_ai/test_gemini_batch_embeddings.py +++ b/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 +