From 4dc277e427a14c437a7a23c4db1f59853af6c287 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Mon, 9 Mar 2026 10:21:29 +0530 Subject: [PATCH] fix(vertex_ai): strip LiteLLM-internal keys from extra_body before merging to Gemini request MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit PR #20950 added extra_body forwarding to Vertex AI Gemini. LiteLLM-internal keys (cache, tags) were being merged into the request body, causing Vertex AI to reject with 400: 'Unknown name "cache": Cannot find field.' - Add _LITELLM_INTERNAL_EXTRA_BODY_KEYS frozenset (cache, tags) - Skip these keys in _pop_and_merge_extra_body before merging - Add regression tests for cache and tags stripping Fixes regression from 1.79.3 → 1.81.12 when using proxy cache with extra_body={"cache": {"use-cache": True, "ttl": 86400}} Made-with: Cursor --- .../llms/vertex_ai/gemini/transformation.py | 6 ++ ...odel_prices_and_context_window_backup.json | 36 ++++++++++ .../test_vertex_ai_gemini_transformation.py | 69 +++++++++++++++++++ 3 files changed, 111 insertions(+) diff --git a/litellm/llms/vertex_ai/gemini/transformation.py b/litellm/llms/vertex_ai/gemini/transformation.py index 57889284a8c..54148c94925 100644 --- a/litellm/llms/vertex_ai/gemini/transformation.py +++ b/litellm/llms/vertex_ai/gemini/transformation.py @@ -529,12 +529,18 @@ def _gemini_convert_messages_with_history( # noqa: PLR0915 raise e +# Keys that LiteLLM consumes internally and must never be forwarded to the +_LITELLM_INTERNAL_EXTRA_BODY_KEYS: frozenset = frozenset({"cache", "tags"}) + + def _pop_and_merge_extra_body(data: RequestBody, optional_params: dict) -> None: """Pop extra_body from optional_params and shallow-merge into data, deep-merging dict values.""" extra_body: Optional[dict] = optional_params.pop("extra_body", None) if extra_body is not None: data_dict: dict = data # type: ignore[assignment] for k, v in extra_body.items(): + if k in _LITELLM_INTERNAL_EXTRA_BODY_KEYS: + continue if k in data_dict and isinstance(data_dict[k], dict) and isinstance(v, dict): data_dict[k].update(v) else: diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 900894f74d6..ae256ed0784 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -16799,6 +16799,42 @@ "supports_vision": true, "supports_web_search": true }, + "gemini/gemini-3.1-flash-image-preview": { + "input_cost_per_token": 2.5e-07, + "input_cost_per_token_batches": 1.25e-07, + "litellm_provider": "gemini", + "max_input_tokens": 65536, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "image_generation", + "output_cost_per_image": 0.045, + "output_cost_per_image_token": 6e-05, + "output_cost_per_image_token_batches": 3e-05, + "output_cost_per_token": 1.5e-06, + "output_cost_per_token_batches": 7.5e-07, + "rpm": 1000, + "tpm": 4000000, + "source": "https://ai.google.dev/gemini-api/docs/pricing#gemini-3.1-flash-image-preview", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text", + "image" + ], + "supports_function_calling": false, + "supports_prompt_caching": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_vision": true, + "supports_web_search": true + }, "gemini/deep-research-pro-preview-12-2025": { "input_cost_per_image": 0.0011, "input_cost_per_token": 2e-06, diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py index b264964b14b..f3c82e439cf 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py @@ -126,6 +126,75 @@ def test_vertex_ai_includes_labels(): +def test_extra_body_cache_not_forwarded_to_vertex_ai(): + """ + 'cache' inside extra_body is a LiteLLM-internal proxy caching control. + It must NOT be forwarded to the Vertex AI request body. + + Regression test for: "Invalid JSON payload received. Unknown name \"cache\": Cannot find field." + Vertex AI enforces a strict JSON schema and rejects any unknown field. + """ + messages = [{"role": "user", "content": "test"}] + optional_params = { + "extra_body": { + "cache": {"use-cache": True, "ttl": 86400}, # LiteLLM-internal + "some_vertex_param": "value", # legitimate provider extra + }, + } + litellm_params = {} + + result = _transform_request_body( + messages=messages, + model="gemini-2.5-pro", + optional_params=optional_params, + custom_llm_provider="vertex_ai", + litellm_params=litellm_params, + cached_content=None, + ) + + # 'cache' must be stripped — Vertex AI has no such field + assert "cache" not in result, ( + "extra_body.cache must not be forwarded to Vertex AI. " + "Vertex AI rejects it with 400: Unknown name \"cache\": Cannot find field." + ) + + # Other legitimate extra_body keys should still pass through + assert "some_vertex_param" in result + assert result["some_vertex_param"] == "value" + + # Core request fields must be present + assert "contents" in result + + +def test_extra_body_tags_not_forwarded_to_vertex_ai(): + """ + 'tags' inside extra_body is a LiteLLM-internal param for logging/tracking. + It must NOT be forwarded to the Vertex AI request body. + Documented in litellm_proxy.md: "Send tags by including them in the extra_body parameter" + """ + messages = [{"role": "user", "content": "test"}] + optional_params = { + "extra_body": { + "tags": ["user:alice", "env:prod"], + "custom_param": "allowed", + }, + } + litellm_params = {} + + result = _transform_request_body( + messages=messages, + model="gemini-2.5-pro", + optional_params=optional_params, + custom_llm_provider="vertex_ai", + litellm_params=litellm_params, + cached_content=None, + ) + + assert "tags" not in result + assert "custom_param" in result + assert result["custom_param"] == "allowed" + + def test_metadata_to_labels_vertex_only(): """Test that metadata->labels conversion only happens for Vertex AI""" messages = [{"role": "user", "content": "test"}]