fix(vertex_ai): strip LiteLLM-internal keys from extra_body before merging to Gemini request

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
This commit is contained in:
Sameer Kankute 2026-03-09 10:21:29 +05:30
parent 28b312f87a
commit 4dc277e427
3 changed files with 111 additions and 0 deletions

View file

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

View file

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

View file

@ -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"}]