Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -384,7 +384,7 @@ def _add_cache_control_if_applicable(
cache_control: Final = (
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)):
if cache_control and model and self.target_consumes_cache_control(model):
# TypedDict objects support dict operations at runtime
# Use type ignore consistent with codebase pattern (see anthropic/chat/transformation.py:432)
if isinstance(target, dict):
Expand Down Expand Up @@ -677,6 +677,10 @@ def is_bedrock_arn_model(model: str) -> bool:
model_lower: Final = model.lower()
return "arn:" in model_lower and ":bedrock:" in model_lower

@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()

@staticmethod
def translate_thinking_for_model(
thinking: AnthropicThinkingParam,
Expand Down
49 changes: 22 additions & 27 deletions litellm/llms/vertex_ai/context_caching/transformation.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@

import re
from collections.abc import Sequence
from types import MappingProxyType
from typing import Final, Literal

from litellm.types.llms.openai import AllMessageValues
Expand Down Expand Up @@ -57,7 +58,7 @@ def extract_ttl_from_cached_messages(messages: list[AllMessageValues]) -> str |
messages: List of messages to extract TTL from

Returns:
Optional[str]: TTL string in format "3600s" or None if not found/invalid
Optional[str]: TTL normalized to Gemini's "<seconds>s" form, or None if not found/invalid
"""
for message in messages:
if not is_cached_message(message):
Expand All @@ -79,40 +80,34 @@ def extract_ttl_from_cached_messages(messages: list[AllMessageValues]) -> str |
if cache_control.get("type") != "ephemeral":
continue

ttl = cache_control.get("ttl")
if ttl and _is_valid_ttl_format(ttl):
return str(ttl)
normalized_ttl = _normalize_ttl_to_seconds(cache_control.get("ttl"))
if normalized_ttl is not None:
return normalized_ttl

return None


def _is_valid_ttl_format(ttl: str) -> bool:
"""
Validate TTL format. Should be a string ending with 's' for seconds.
Examples: "3600s", "7200s", "1.5s"
_TTL_PATTERN: Final = re.compile(r"^([0-9]*\.?[0-9]+)([smh])$")
_TTL_UNIT_SECONDS: Final = MappingProxyType({"s": 1, "m": 60, "h": 3600})

Args:
ttl: TTL string to validate

Returns:
bool: True if valid format, False otherwise
def _normalize_ttl_to_seconds(ttl: object) -> str | None:
"""
Gemini's cachedContents API only takes a TTL as "<seconds>s", while Anthropic clients
(Claude Code among them) send the minute and hour units the Anthropic API defines, "5m"
and "1h". Returns the Gemini form for any of the three units, or None for a missing,
non-positive, or unparseable value so the cache falls back to Gemini's default TTL.
"""
if not isinstance(ttl, str):
return False

# TTL should end with 's' and contain a valid number before it
pattern: Final = r"^([0-9]*\.?[0-9]+)s$"
match: Final = re.match(pattern, ttl)

if not match:
return False

try:
# Ensure the numeric part is valid and positive
numeric_part: Final = float(match.group(1))
return numeric_part > 0
except ValueError:
return False
return None
match: Final = _TTL_PATTERN.match(ttl)
if match is None:
return None
value: Final = float(match.group(1))
if value <= 0:
return None
seconds: Final = round(value * _TTL_UNIT_SECONDS[match.group(2)], 9)
return f"{seconds:.9f}".rstrip("0").rstrip(".") + "s"


def separate_cached_messages(
Expand Down
2 changes: 2 additions & 0 deletions litellm/model_prices_and_context_window_backup.json
Original file line number Diff line number Diff line change
Expand Up @@ -25925,6 +25925,7 @@
"supports_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"prompt_cache_min_tokens": 2048,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_system_messages": true,
Expand Down Expand Up @@ -27895,6 +27896,7 @@
"supports_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"prompt_cache_min_tokens": 2048,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_system_messages": true,
Expand Down
2 changes: 2 additions & 0 deletions model_prices_and_context_window.json
Original file line number Diff line number Diff line change
Expand Up @@ -25925,6 +25925,7 @@
"supports_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"prompt_cache_min_tokens": 2048,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_system_messages": true,
Expand Down Expand Up @@ -27895,6 +27896,7 @@
"supports_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"prompt_cache_min_tokens": 2048,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_system_messages": true,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -2102,14 +2102,17 @@ def test_should_add_cache_control_for_anthropic_model():


def test_should_not_add_cache_control_for_non_anthropic_model():
"""Should not add cache_control for non-Anthropic models."""
"""Should not add cache_control for providers that reject an explicit cache_control field.

OpenAI/Azure do prompt caching implicitly and 400 on an unexpected
cache_control field, so it must not be forwarded to them.
"""
adapter = LiteLLMAnthropicMessagesAdapter()
cache_control = {"type": "ephemeral"}

for model in [
CACHE_CONTROL_NON_ANTHROPIC_MODEL,
"openai/gpt-4-turbo",
"gemini-pro",
]:
target = {}
adapter._add_cache_control_if_applicable(
Expand All @@ -2118,6 +2121,54 @@ def test_should_not_add_cache_control_for_non_anthropic_model():
assert "cache_control" not in target


def test_should_add_cache_control_for_gemini_model():
"""Should add cache_control for Gemini / Vertex Gemini targets.

These consume anthropic-style cache_control blocks via the Gemini context
caching path, so /v1/messages requests (e.g. Claude Code) routed to a
Gemini model must keep it. Regression for the adapter dropping the field
before it reaches the Gemini transformation.
"""
adapter = LiteLLMAnthropicMessagesAdapter()
cache_control = {"type": "ephemeral", "ttl": "1h"}

for model in [
"gemini-3.5-flash",
"gemini/gemini-3.5-flash",
"gemini-3.1-pro-preview",
"vertex_ai/gemini-2.5-pro",
]:
target = {}
adapter._add_cache_control_if_applicable(
{"cache_control": cache_control}, target, model
)
assert target.get("cache_control") == cache_control


def test_cache_control_preserved_in_text_content_for_gemini():
"""cache_control must survive message translation for a Gemini target."""
anthropic_messages = [
AnthropicMessagesUserMessageParam(
role="user",
content=[
{
"type": "text",
"text": "This is cached content",
"cache_control": {"type": "ephemeral", "ttl": "1h"},
}
],
)
]

adapter = LiteLLMAnthropicMessagesAdapter()
result = adapter.translate_anthropic_messages_to_openai(
messages=anthropic_messages, model="gemini/gemini-3.5-flash"
)

assert len(result) == 1
assert result[0]["content"][0]["cache_control"] == {"type": "ephemeral", "ttl": "1h"}


def test_should_not_add_cache_control_when_none():
"""Should not add cache_control when source has None or empty cache_control."""
adapter = LiteLLMAnthropicMessagesAdapter()
Expand Down
Original file line number Diff line number Diff line change
@@ -1,45 +1,77 @@
import pytest
from litellm.llms.vertex_ai.context_caching.transformation import (
_normalize_ttl_to_seconds,
extract_ttl_from_cached_messages,
_is_valid_ttl_format,
transform_openai_messages_to_gemini_context_caching,
)


class TestTTLValidation:
"""Test TTL format validation"""

def test_valid_ttl_formats(self):
"""Test various valid TTL formats"""
valid_ttls = ["3600s", "1s", "7200s", "1.5s", "0.1s", "86400s", "123.456s"]

for ttl in valid_ttls:
assert _is_valid_ttl_format(ttl), f"TTL {ttl} should be valid"

def test_invalid_ttl_formats(self):
"""Test various invalid TTL formats"""
invalid_ttls = [
"3600", # missing 's'
"s", # missing number
"-1s", # negative number
"0s", # zero
"3600m", # wrong unit
"abc.s", # invalid number
"", # empty string
"3600.s", # invalid decimal
"3600 s", # space
"3600ss", # extra 's'
None, # None
123, # not a string
]
class TestTTLNormalization:
"""Gemini only takes "<seconds>s"; Anthropic clients send "5m" and "1h" too"""

for ttl in invalid_ttls:
assert not _is_valid_ttl_format(ttl), f"TTL {ttl} should be invalid"
@pytest.mark.parametrize(
"ttl, expected",
[
("3600s", "3600s"),
("1s", "1s"),
("1.5s", "1.5s"),
("0.1s", "0.1s"),
("123.456s", "123.456s"),
("1.3333333333333333s", "1.333333333s"),
("5m", "300s"),
("90m", "5400s"),
("1h", "3600s"),
("0.5h", "1800s"),
("48h", "172800s"),
],
)
def test_normalizes_supported_units_to_seconds(self, ttl, expected):
assert _normalize_ttl_to_seconds(ttl) == expected

@pytest.mark.parametrize(
"ttl",
[
"3600",
"s",
"-1s",
"0s",
"0m",
"0h",
"5d",
"abc.s",
"",
"3600.s",
"3600 s",
"3600ss",
"1 h",
None,
123,
],
)
def test_rejects_unparseable_ttl(self, ttl):
assert _normalize_ttl_to_seconds(ttl) is None


class TestTTLExtraction:
"""Test TTL extraction from cached messages"""

@pytest.mark.parametrize("ttl, expected", [("1h", "3600s"), ("5m", "300s")])
def test_extract_ttl_normalizes_anthropic_units(self, ttl, expected):
messages = [
{
"role": "system",
"content": [
{
"type": "text",
"text": "cached",
"cache_control": {"type": "ephemeral", "ttl": ttl},
}
],
}
]

assert extract_ttl_from_cached_messages(messages) == expected

def test_extract_ttl_from_single_message(self):
"""Test extracting TTL from a single cached message"""
messages = [
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -1396,6 +1396,48 @@ def test_check_and_create_cache_skips_when_below_min_tokens(
# Restart the patcher so teardown_method can stop it cleanly
self._token_check_patcher.start()

def test_check_and_create_cache_skips_between_default_and_gemini_2_5_pro_minimum(
self, local_model_cost_map
):
"""Gemini 2.5 Pro needs 2048 cached tokens, twice the provider-agnostic default.

Content between the two used to reach Google's cachedContents endpoint and 400
with "Cached content is too small".
"""
model = "gemini-2.5-pro"
self._token_check_patcher.stop()

cached_messages = [
{
"role": "system",
"content": " ".join(["word"] * 1500),
"cache_control": {"type": "ephemeral"},
}
]
non_cached_messages = [{"role": "user", "content": "Hello"}]

messages, _, returned_cache = 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 messages == cached_messages + non_cached_messages
assert returned_cache is None
self.mock_client.post.assert_not_called()

self._token_check_patcher.start()

@pytest.mark.parametrize(
"custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]
)
Expand Down
Loading