This commit is contained in:
Scripts May 2026-11-06 04:59:21 +00:00 • committed by GitHub
commit 00a394dfbd
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 272 additions and 9 deletions

View file

@ -743,7 +743,22 @@ class LiteLLMAnthropicMessagesAdapter:
@classmethod
def target_consumes_cache_control(cls, model: str) -> bool:
return cls.is_anthropic_claude_model(model) or cls.is_bedrock_arn_model(model) or "gemini" in model.lower()
if cls.is_anthropic_claude_model(model) or cls.is_bedrock_arn_model(model):
return True
if "gemini" in model.lower():
from litellm.utils import is_explicitly_disabled_factory
if "/" not in model:
if (
is_explicitly_disabled_factory(f"gemini/{model}", None, "supports_prompt_caching")
or is_explicitly_disabled_factory(f"vertex_ai/{model}", None, "supports_prompt_caching")
):
return False
if is_explicitly_disabled_factory(model, None, "supports_prompt_caching"):
return False
return True
return False
@staticmethod
def translate_thinking_for_model(

View file

@ -130,6 +130,14 @@ def _context_caching_skip_reason(
cache_model: str,
custom_llm_provider: Literal["vertex_ai", "vertex_ai_beta", "gemini"],
) -> str | None:
from litellm.utils import is_explicitly_disabled_factory
if is_explicitly_disabled_factory(deployment_model, custom_llm_provider, "supports_prompt_caching") or (
cache_model != deployment_model
and is_explicitly_disabled_factory(cache_model, custom_llm_provider, "supports_prompt_caching")
):
return f"Vertex AI context caching: model {deployment_model} does not support explicit context caching."
if not cached_messages_end_on_supported_turn(cached_messages):
return (
"Vertex AI context caching: cached message block ends on a model turn once "
@ -499,9 +507,9 @@ class ContextCachingEndpoints(VertexBase):
except _CachedContentsNotServedError:
verbose_logger.debug("Gemini context caching: %s answers 404. Skipping context caching.", url)
return messages, optional_params, None
for tool_param in _TOOL_PARAMS:
optional_params.pop(tool_param, None)
if google_cache_name:
for tool_param in _TOOL_PARAMS:
optional_params.pop(tool_param, None)
return non_cached_messages, optional_params, google_cache_name
## TRANSFORM REQUEST
@ -537,6 +545,17 @@ class ContextCachingEndpoints(VertexBase):
)
response.raise_for_status()
except httpx.HTTPStatusError as err:
if "does not support explicit context caching" in err.response.text or (
err.response.status_code == 400
and "explicit" in err.response.text.lower()
and "caching" in err.response.text.lower()
):
verbose_logger.warning(
"Model %s does not support explicit context caching (%s). Falling back to non-cached request.",
model,
err.response.text,
)
return messages, optional_params, None
error_code: Final = err.response.status_code
raise VertexAIError(status_code=error_code, message=err.response.text)
except httpx.TimeoutException:
@ -552,6 +571,8 @@ class ContextCachingEndpoints(VertexBase):
raw_response_cached.get("name"),
raw_response_cached.get("expireTime"),
)
for tool_param in _TOOL_PARAMS:
optional_params.pop(tool_param, None)
return (
non_cached_messages,
optional_params,
@ -689,6 +710,17 @@ class ContextCachingEndpoints(VertexBase):
)
response.raise_for_status()
except httpx.HTTPStatusError as err:
if "does not support explicit context caching" in err.response.text or (
err.response.status_code == 400
and "explicit" in err.response.text.lower()
and "caching" in err.response.text.lower()
):
verbose_logger.warning(
"Model %s does not support explicit context caching (%s). Falling back to non-cached request.",
model,
err.response.text,
)
return None
error_code: Final = err.response.status_code
raise VertexAIError(status_code=error_code, message=err.response.text)
except httpx.TimeoutException:
@ -708,6 +740,10 @@ class ContextCachingEndpoints(VertexBase):
except _CachedContentsNotServedError:
verbose_logger.debug("Gemini context caching: %s answers 404. Skipping context caching.", url)
return messages, optional_params, None
if cache_name is None:
return messages, optional_params, None
for tool_param in _TOOL_PARAMS:
optional_params.pop(tool_param, None)
return non_cached_messages, optional_params, cache_name

View file

@ -1262,6 +1262,11 @@ def sync_transform_request_body(
context_caching_endpoints: Final = ContextCachingEndpoints()
if context_caching:
model_info = litellm_params.get("model_info") or {}
if model_info.get("supports_prompt_caching") is False:
context_caching = False
(
messages,
optional_params,
@ -1398,6 +1403,11 @@ async def async_transform_request_body(
context_caching_endpoints: Final = ContextCachingEndpoints()
if context_caching:
model_info = litellm_params.get("model_info") or {}
if model_info.get("supports_prompt_caching") is False:
context_caching = False
(
messages,
optional_params,

View file

@ -78192,7 +78192,7 @@
"supports_native_streaming": true,
"supports_parallel_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_prompt_caching": false,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_system_messages": true,
@ -78252,7 +78252,7 @@
"supports_native_streaming": true,
"supports_parallel_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_prompt_caching": false,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_system_messages": true,
@ -78310,7 +78310,7 @@
"supports_native_streaming": true,
"supports_parallel_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_prompt_caching": false,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_system_messages": true,

View file

@ -78192,7 +78192,7 @@
"supports_native_streaming": true,
"supports_parallel_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_prompt_caching": false,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_system_messages": true,
@ -78252,7 +78252,7 @@
"supports_native_streaming": true,
"supports_parallel_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_prompt_caching": false,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_system_messages": true,
@ -78310,7 +78310,7 @@
"supports_native_streaming": true,
"supports_parallel_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_prompt_caching": false,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_system_messages": true,

View file

@ -0,0 +1,202 @@
import json
from unittest.mock import AsyncMock, MagicMock
import httpx
import pytest
from litellm.llms.anthropic.messages.adapters.transformation import (
LiteLLMAnthropicMessagesAdapter,
)
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
from litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching import (
ContextCachingEndpoints,
)
from litellm.llms.vertex_ai.gemini.transformation import (
sync_transform_request_body,
)
def test_target_consumes_cache_control_skips_when_model_disables_prompt_caching(local_model_cost_map):
"""LiteLLMAnthropicMessagesAdapter should not consume cache_control for models explicitly disabling prompt caching."""
assert LiteLLMAnthropicMessagesAdapter.target_consumes_cache_control("claude-3-5-sonnet-20241022") is True
assert LiteLLMAnthropicMessagesAdapter.target_consumes_cache_control("gemini/gemini-2.5-flash") is True
assert LiteLLMAnthropicMessagesAdapter.target_consumes_cache_control("gemini/gemini-4-argon") is False
assert LiteLLMAnthropicMessagesAdapter.target_consumes_cache_control("vertex_ai/gemini-4-argon") is False
# Bare model name without provider prefix should also be recognized
assert LiteLLMAnthropicMessagesAdapter.target_consumes_cache_control("gemini-4-argon") is False
def test_check_and_create_cache_skips_when_model_disables_prompt_caching(local_model_cost_map):
"""ContextCachingEndpoints should skip explicit cache creation when model has supports_prompt_caching: false."""
endpoints = ContextCachingEndpoints()
messages = [
{
"role": "user",
"content": [
{
"type": "text",
"text": "hello world test token count " * 1000,
"cache_control": {"type": "ephemeral"},
}
],
}
]
optional_params = {"tools": [{"type": "function", "function": {"name": "test"}}]}
res_messages, res_params, cache_id = endpoints.check_and_create_cache(
messages=messages,
optional_params=optional_params,
api_key="test-key",
api_base=None,
model="gemini/gemini-4-argon",
client=MagicMock(),
timeout=10,
logging_obj=MagicMock(),
custom_llm_provider="gemini",
vertex_project=None,
vertex_location=None,
vertex_auth_header=None,
)
assert cache_id is None
# All messages and tools should remain intact
assert len(res_messages) == len(messages)
assert "tools" in res_params
@pytest.mark.asyncio
async def test_async_check_and_create_cache_fallback_on_400_explicit_caching_unsupported(local_model_cost_map):
"""When Vertex AI returns 400 'does not support explicit context caching', gracefully fallback to non-cached request."""
endpoints = ContextCachingEndpoints()
endpoints.async_check_cache = AsyncMock(return_value=None)
messages = [
{
"role": "user",
"content": [
{
"type": "text",
"text": "hello world test token count " * 1000,
"cache_control": {"type": "ephemeral"},
}
],
}
]
optional_params = {"tools": [{"type": "function", "function": {"name": "test"}}]}
mock_client = MagicMock(spec=AsyncHTTPHandler)
mock_response = MagicMock()
mock_response.status_code = 400
mock_response.text = json.dumps({
"error": {
"code": 400,
"message": "Model gemini-4-argon does not support explicit context caching.",
"status": "INVALID_ARGUMENT",
}
})
mock_client.post = AsyncMock(
side_effect=httpx.HTTPStatusError("Bad Request", request=MagicMock(), response=mock_response)
)
res_messages, res_params, cache_id = await endpoints.async_check_and_create_cache(
messages=messages,
optional_params=optional_params,
api_key="test-key",
api_base=None,
model="gemini/gemini-2.5-flash",
client=mock_client,
timeout=10,
logging_obj=MagicMock(),
custom_llm_provider="gemini",
vertex_project=None,
vertex_location=None,
vertex_auth_header=None,
)
assert cache_id is None
assert mock_client.post.called
# Critical checks: messages must not be dropped, tools must not be stripped!
assert len(res_messages) == len(messages)
assert "tools" in res_params
def test_sync_check_and_create_cache_fallback_on_400_explicit_caching_unsupported(local_model_cost_map):
"""Sync check_and_create_cache fallback on 400 explicit caching unsupported."""
endpoints = ContextCachingEndpoints()
endpoints.check_cache = MagicMock(return_value=None)
messages = [
{
"role": "user",
"content": [
{
"type": "text",
"text": "hello world test token count " * 1000,
"cache_control": {"type": "ephemeral"},
}
],
}
]
optional_params = {"tools": [{"type": "function", "function": {"name": "test"}}]}
mock_client = MagicMock(spec=HTTPHandler)
mock_response = MagicMock()
mock_response.status_code = 400
mock_response.text = json.dumps({
"error": {
"code": 400,
"message": "Model gemini-4-argon does not support explicit context caching.",
"status": "INVALID_ARGUMENT",
}
})
mock_client.post.side_effect = httpx.HTTPStatusError(
"Bad Request", request=MagicMock(), response=mock_response
)
res_messages, res_params, cache_id = endpoints.check_and_create_cache(
messages=messages,
optional_params=optional_params,
api_key="test-key",
api_base=None,
model="gemini/gemini-2.5-flash",
client=mock_client,
timeout=10,
logging_obj=MagicMock(),
custom_llm_provider="gemini",
vertex_project=None,
vertex_location=None,
vertex_auth_header=None,
)
assert cache_id is None
assert mock_client.post.called
# Critical checks: messages must not be dropped, tools must not be stripped!
assert len(res_messages) == len(messages)
assert "tools" in res_params
def test_sync_transform_request_body_skips_caching_on_model_info_override(monkeypatch):
"""sync_transform_request_body should respect supports_prompt_caching: false in model_info."""
mock_endpoints = MagicMock()
mock_endpoints.check_and_create_cache.return_value = ([], {}, None)
monkeypatch.setattr(
"litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.ContextCachingEndpoints",
MagicMock(return_value=mock_endpoints),
)
messages = [{"role": "user", "content": "hello"}]
sync_transform_request_body(
messages=messages,
model="gemini-2.5-flash",
custom_llm_provider="gemini",
litellm_params={"model_info": {"supports_prompt_caching": False}},
optional_params={},
gemini_api_key="test-key",
api_base=None,
client=None,
timeout=None,
extra_headers=None,
logging_obj=MagicMock(),
vertex_project=None,
vertex_location=None,
vertex_auth_header=None,
)
# context_caching=False must be passed to check_and_create_cache
assert mock_endpoints.check_and_create_cache.called
_, kwargs = mock_endpoints.check_and_create_cache.call_args
assert kwargs.get("context_caching") is False