feat(prompt-caching): match affinity TTL to cache control TTL

Add logic to extract TTL from cache_control blocks and use 1 hour (3600s) for "1h" TTL or default 5 minutes (300s) otherwise. Apply the extracted TTL when storing model affinity in both sync and async add_model_id methods.
This commit is contained in:
nuernber 2026-09-11 12:19:08 -07:00
parent acb9086f29
commit 08f348bebd
2 changed files with 68 additions and 2 deletions

View file

@ -139,6 +139,27 @@ class PromptCachingCache:
return cacheable_prefix
@staticmethod
def get_prompt_caching_ttl(messages: list[AllMessageValues] | None) -> int:
if messages is None:
return 300
cacheable_prefix: Final = PromptCachingCache.extract_cacheable_prefix(messages)
cache_control_ttls: Final = tuple(
cache_control.get("ttl")
for message in cacheable_prefix
for cache_control in (
message.get("cache_control"),
*(
content_block.get("cache_control")
for content_block in (message.get("content") if isinstance(message.get("content"), list) else ())
if isinstance(content_block, dict)
),
)
if isinstance(cache_control, dict) and cache_control.get("type") == "ephemeral"
)
return 3600 if "1h" in cache_control_ttls else 300
@staticmethod
def get_prompt_caching_cache_key(
messages: list[AllMessageValues] | None,
@ -189,7 +210,11 @@ class PromptCachingCache:
if cache_key is None:
return
self.cache.set_cache(cache_key, PromptCachingCacheValue(model_id=model_id), ttl=300)
self.cache.set_cache(
cache_key,
PromptCachingCacheValue(model_id=model_id),
ttl=PromptCachingCache.get_prompt_caching_ttl(messages),
)
return
async def async_add_model_id(
@ -209,7 +234,7 @@ class PromptCachingCache:
await self.cache.async_set_cache(
cache_key,
PromptCachingCacheValue(model_id=model_id),
ttl=300, # store for 5 minutes
ttl=PromptCachingCache.get_prompt_caching_ttl(messages),
)
return

View file

@ -1,6 +1,7 @@
import asyncio
import copy
from typing import List, cast
from unittest.mock import AsyncMock
import pytest
@ -60,6 +61,46 @@ def _messages(word_count: int) -> List[AllMessageValues]:
)
@pytest.mark.parametrize(
("ttl", "expected_affinity_ttl"),
((None, 300), ("5m", 300), ("1h", 3600)),
)
def test_prompt_caching_affinity_ttl_matches_cache_control(ttl: str | None, expected_affinity_ttl: int):
cache_control: dict[str, str] = {"type": "ephemeral", **({"ttl": ttl} if ttl is not None else {})}
messages = cast(
List[AllMessageValues],
[{"role": "system", "content": [{"type": "text", "text": "cached", "cache_control": cache_control}]}],
)
assert PromptCachingCache.get_prompt_caching_ttl(messages) == expected_affinity_ttl
@pytest.mark.asyncio
async def test_async_add_model_id_uses_one_hour_affinity_ttl():
cache = DualCache()
async_set_cache = AsyncMock()
cache.async_set_cache = async_set_cache
messages = cast(
List[AllMessageValues],
[
{
"role": "system",
"content": [
{
"type": "text",
"text": "cached",
"cache_control": {"type": "ephemeral", "ttl": "1h"},
}
],
}
],
)
await PromptCachingCache(cache=cache).async_add_model_id("dep-1", messages, None)
assert async_set_cache.call_args.kwargs["ttl"] == 3600
def test_get_min_token_count_for_deployments_takes_min_across_mixed_group():
"""
A group may legally mix models whose real minimums differ, and one gate decides for every