scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-25 01:02:15 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 80

feat(caching): enhance Gemini context caching by enforcing minimum token requirements to prevent 400s

Browse source
This commit is contained in:
Elliott de Launay 2026-07-09 22:53:04 -04:00
parent 5822aa87ee
commit 0695f0702d
4 changed files with 81 additions and 11 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

12
litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py
View file

@ -311,10 +311,14 @@ class LiteLLMAnthropicMessagesAdapter:
cache_control = (
source.get("cache_control") if isinstance(source, dict) else getattr(source, "cache_control", None)
)
if cache_control and model and (
self.is_anthropic_claude_model(model)
or self.is_bedrock_arn_model(model)
or _is_gemini_model(model, None)
if (
cache_control
and model
and (
self.is_anthropic_claude_model(model)
or self.is_bedrock_arn_model(model)
or _is_gemini_model(model, None)
)
):
# TypedDict objects support dict operations at runtime
# Use type ignore consistent with codebase pattern (see anthropic/chat/transformation.py:432)

20
litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py
View file

@ -4,7 +4,6 @@ import httpx
import litellm
from litellm.caching.caching import Cache, LiteLLMCacheType
from litellm.constants import MINIMUM_PROMPT_CACHE_TOKEN_COUNT
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.llms.custom_httpx.http_handler import (
AsyncHTTPHandler,
@ -22,6 +21,7 @@ from litellm.types.llms.vertex_ai import (
from ..common_utils import VertexAIError, get_vertex_base_url
from ..vertex_llm_base import VertexBase
from .transformation import (
get_gemini_context_caching_min_tokens,
separate_cached_messages,
transform_openai_messages_to_gemini_context_caching,
)
@ -308,17 +308,20 @@ class ContextCachingEndpoints(VertexBase):
if len(cached_messages) == 0:
return messages, optional_params, None
# Gemini requires a minimum of 1024 tokens for context caching.
# Skip caching if the cached content is too small to avoid API errors.
# Gemini's explicit context caching minimum varies by model; creating a
# cache below it returns a 400. Skip caching when the cached content is
# too small to avoid the error.
min_token_count = get_gemini_context_caching_min_tokens(model)
if not is_prompt_caching_valid_prompt(
model=model,
messages=cached_messages,
custom_llm_provider=custom_llm_provider,
min_token_count=min_token_count,
):
verbose_logger.debug(
"Vertex AI context caching: cached content is below minimum token "
"count (%d). Skipping context caching.",
MINIMUM_PROMPT_CACHE_TOKEN_COUNT,
min_token_count,
)
return messages, optional_params, None
@ -459,17 +462,20 @@ class ContextCachingEndpoints(VertexBase):
if len(cached_messages) == 0:
return messages, optional_params, None
# Gemini requires a minimum of 1024 tokens for context caching.
# Skip caching if the cached content is too small to avoid API errors.
# Gemini's explicit context caching minimum varies by model; creating a
# cache below it returns a 400. Skip caching when the cached content is
# too small to avoid the error.
min_token_count = get_gemini_context_caching_min_tokens(model)
if not is_prompt_caching_valid_prompt(
model=model,
messages=cached_messages,
custom_llm_provider=custom_llm_provider,
min_token_count=min_token_count,
):
verbose_logger.debug(
"Vertex AI context caching: cached content is below minimum token "
"count (%d). Skipping context caching.",
MINIMUM_PROMPT_CACHE_TOKEN_COUNT,
min_token_count,
)
return messages, optional_params, None

2
litellm/utils.py
View file

@ -9097,6 +9097,8 @@ def is_prompt_caching_valid_prompt(
nothing here and would silently fall back to the default.
OpenAI's minimum is a flat 1024 across models, which the default already covers.
Pass min_token_count to override this for providers with a different floor
(e.g. Gemini, whose explicit context caching minimum varies by model).
"""
try:
if messages is None and tools is None:

58
tests/test_litellm/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py
View file

@ -1401,6 +1401,64 @@ class TestContextCachingEndpoints:
# Restart the patcher so teardown_method can stop it cleanly
self._token_check_patcher.start()
@pytest.mark.parametrize(
"model, expected_min",
[
("gemini-3.5-flash", 4096),
("gemini/gemini-3.5-flash", 4096),
("gemini-3.1-pro-preview", 4096),
("gemini-2.5-flash", 2048),
("gemini-2.5-pro", 2048),
],
)
@patch(
"litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages"
)
def test_check_and_create_cache_uses_model_specific_min_tokens(
self, mock_separate, model, expected_min
):
"""The Gemini per-model floor must be forwarded to the token-count guard.
A flat 1024 floor let content between 1024 and the real minimum (2048 for
2.5, 4096 for 3.x) reach Gemini and 400. Assert the model-derived floor is
passed so the guard skips instead of erroring.
"""
self._token_check_patcher.stop()
cached_messages = [
{
"role": "system",
"content": "cached",
"cache_control": {"type": "ephemeral"},
}
]
non_cached_messages = [{"role": "user", "content": "Hello"}]
mock_separate.return_value = (cached_messages, non_cached_messages)
with patch(
"litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.is_prompt_caching_valid_prompt",
return_value=False,
) as mock_valid:
self.context_caching.check_and_create_cache(
messages=cached_messages + non_cached_messages,
optional_params=self.sample_optional_params.copy(),
api_key="test_key",
api_base=None,
model=model,
client=self.mock_client,
timeout=30.0,
logging_obj=self.mock_logging,
cached_content=None,
custom_llm_provider="gemini",
vertex_project="test_project",
vertex_location="us-central1",
vertex_auth_header="test_token",
)
assert mock_valid.call_args.kwargs["min_token_count"] == expected_min
self._token_check_patcher.start()
@pytest.mark.parametrize(
"custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]
)