mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-07 02:59:05 +00:00
fix(vertex): drop extra_body.cache from Gemini request payload
This commit is contained in:
parent
11e2bdbaec
commit
2672294c01
2 changed files with 49 additions and 0 deletions
|
|
@ -533,6 +533,7 @@ 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:
|
||||
extra_body.pop("cache", None)
|
||||
data_dict: dict = data # type: ignore[assignment]
|
||||
for k, v in extra_body.items():
|
||||
if k in data_dict and isinstance(data_dict[k], dict) and isinstance(v, dict):
|
||||
|
|
|
|||
|
|
@ -125,6 +125,54 @@ def test_vertex_ai_includes_labels():
|
|||
assert result["labels"] == {"project": "test", "team": "ai"}
|
||||
|
||||
|
||||
def test_vertex_ai_filters_cache_from_extra_body():
|
||||
"""Test that LiteLLM cache controls in extra_body are not forwarded to Vertex payloads."""
|
||||
messages = [{"role": "user", "content": "test"}]
|
||||
optional_params = {
|
||||
"extra_body": {
|
||||
"cache": {"use-cache": True, "ttl": 86400},
|
||||
"labels": {"project": "test"},
|
||||
}
|
||||
}
|
||||
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 "cache" not in result
|
||||
assert result.get("labels") == {"project": "test"}
|
||||
|
||||
|
||||
def test_gemini_filters_cache_from_extra_body():
|
||||
"""Test that LiteLLM cache controls in extra_body are not forwarded to Gemini payloads."""
|
||||
messages = [{"role": "user", "content": "test"}]
|
||||
optional_params = {
|
||||
"extra_body": {
|
||||
"cache": {"use-cache": True, "ttl": 86400},
|
||||
"generationConfig": {"temperature": 0.2},
|
||||
}
|
||||
}
|
||||
litellm_params = {}
|
||||
|
||||
result = _transform_request_body(
|
||||
messages=messages,
|
||||
model="gemini-2.5-pro",
|
||||
optional_params=optional_params,
|
||||
custom_llm_provider="gemini",
|
||||
litellm_params=litellm_params,
|
||||
cached_content=None,
|
||||
)
|
||||
|
||||
assert "cache" not in result
|
||||
assert result["generationConfig"]["temperature"] == 0.2
|
||||
|
||||
|
||||
|
||||
def test_metadata_to_labels_vertex_only():
|
||||
"""Test that metadata->labels conversion only happens for Vertex AI"""
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue