mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-14 23:21:35 +00:00
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:
parent
acb9086f29
commit
08f348bebd
2 changed files with 68 additions and 2 deletions
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue