scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-29 06:33:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 267

fix(cache): propagate unavailable cache keys safely

Browse source
This commit is contained in:
Maximilian Rink 2026-10-10 18:20:15 +02:00
parent 877981dc4c
commit c3ff7cd320
5 changed files with 77 additions and 37 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

17
litellm/caching/caching.py
View file

@ -451,7 +451,8 @@ class Cache:
if litellm.enable_caching_on_provider_specific_optional_params is True: # feature flagged for now
if kwargs[param] is None:
continue # ignore None params
cache_key += f"{param}: {kwargs[param]}"
optional_param_value: object = kwargs[param]
cache_key += f"{param}: {optional_param_value}"
if is_semantic_cache:
cache_key += self._get_semantic_cache_tenant_scope(kwargs)
@ -690,9 +691,7 @@ class Cache:
try:
from litellm.proxy._types import UserAPIKeyAuth
except ImportError:
class UserAPIKeyAuth: # pragma: no cover - proxy types are available in proxy deployments
pass
return False, None
for metadata in metadata_sources:
auth_object: object | None = metadata.get("user_api_key_auth")
@ -729,6 +728,14 @@ class Cache:
return None
return self._add_namespace_to_cache_key(explicit_key, **kwargs)
def prepare_explicit_cache_key(
self,
explicit_key: object,
**kwargs: object, # kwargs-ok: auth metadata scopes caller keys
) -> str | None:
"""Normalize a caller-provided key for integrations outside this class."""
return self._prepare_explicit_cache_key(explicit_key, **kwargs)
def generate_streaming_content(self, content):
chunk_size: Final = 5 # Adjust the chunk size as needed
for i in range(0, len(content), chunk_size):
@ -1047,7 +1054,7 @@ class Cache:
responses distribute it evenly (with remainder) so that summing all
per-item values on retrieval reconstructs the original total.
"""
if result.usage is None or result.usage.prompt_tokens is None:
if result.usage is None:
return None
total: Final = result.usage.prompt_tokens

25
litellm/caching/caching_handler.py
View file

@ -19,7 +19,7 @@ import datetime
import inspect
import time
from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable, Generator, Mapping
from typing import TYPE_CHECKING, Any, Final, Optional, TypeVar
from typing import TYPE_CHECKING, Any, Final, Optional, TypeVar, cast
from pydantic import ConfigDict, SkipValidation, TypeAdapter, ValidationError
@ -406,7 +406,6 @@ class LLMCachingHandler:
call_type == CallTypes.aembedding.value
and cached_result is not None
and isinstance(cached_result, list)
and litellm.cache is not None
and not isinstance(litellm.cache.cache, S3Cache) # s3 doesn't support bulk writing. Exclude.
):
(
@ -465,7 +464,9 @@ class LLMCachingHandler:
if new_kwargs.get("metadata") is None:
new_kwargs.pop("metadata", None)
if new_kwargs.get("stream") is True and "cache_key" not in new_kwargs:
new_kwargs["cache_key"] = litellm.cache.get_cache_key(**new_kwargs)
derived_cache_key = litellm.cache.get_cache_key(**cast(dict[str, object], new_kwargs))
if derived_cache_key is not None:
new_kwargs["cache_key"] = derived_cache_key
self.request_kwargs = _drop_logging_obj_from_kwargs(new_kwargs)
print_verbose("Checking Sync Cache")
logging_obj.reset_caching_details()
@ -886,13 +887,17 @@ class LLMCachingHandler:
if _is_response_cache_excluded(model=model, kwargs=new_kwargs):
return None
if new_kwargs.get("stream") is True and "cache_key" not in new_kwargs:
new_kwargs["cache_key"] = litellm.cache.get_cache_key(**new_kwargs)
derived_cache_key = litellm.cache.get_cache_key(**new_kwargs)
if derived_cache_key is not None:
new_kwargs["cache_key"] = derived_cache_key
cached_result: object | None = None
if call_type == CallTypes.aembedding.value:
new_kwargs["input"] = self.handle_kwargs_input_list_or_str(new_kwargs)
tasks: Final[list[Awaitable[object]]] = []
for idx, i in enumerate(new_kwargs["input"]):
preset_cache_key = litellm.cache.get_cache_key(**{**new_kwargs, "input": i})
preset_cache_key = litellm.cache.get_cache_key(**cast(dict[str, object], {**new_kwargs, "input": i}))
if preset_cache_key is None:
return None
tasks.append(
litellm.cache.async_get_cache(
cache_key=preset_cache_key,
@ -903,7 +908,7 @@ class LLMCachingHandler:
entries: Final = await asyncio.gather(*tasks)
cached_result = [_current_format_embedding_entry(entry) for entry in entries]
## check if cached result is None ##
if cached_result is not None and isinstance(cached_result, list):
if isinstance(cached_result, list):
# set cached_result to None if all elements are None
if all(result is None for result in cached_result):
cached_result = None
@ -911,9 +916,13 @@ class LLMCachingHandler:
request_kwargs: Final = new_kwargs.copy()
request_cache_key: Final = _request_cache_key(request_kwargs)
request_kwargs.pop("cache_key", None)
derived_cache_key = litellm.cache.get_cache_key(**request_kwargs)
if request_cache_key is None and derived_cache_key is None:
self.preset_cache_key = None
return None
if litellm.cache.supports_async() is True:
## check if dual cache is supported ##
self.preset_cache_key = request_cache_key or litellm.cache.get_cache_key(**request_kwargs)
self.preset_cache_key = request_cache_key or derived_cache_key
with response_cache_phase("get"):
cached_result = await litellm.cache.async_get_cache(
dynamic_cache_object=self.dual_cache,
@ -921,7 +930,7 @@ class LLMCachingHandler:
**request_kwargs,
)
else: # fallback for caches that don't support async
self.preset_cache_key = request_cache_key or litellm.cache.get_cache_key(**request_kwargs)
self.preset_cache_key = request_cache_key or derived_cache_key
with response_cache_phase("get"):
cached_result = litellm.cache.get_cache(
dynamic_cache_object=self.dual_cache,

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

@ -1,5 +1,5 @@
from collections.abc import Sequence
from typing import Final, Literal
from typing import Final, Literal, cast
import httpx
from pydantic import ConfigDict, RootModel, with_config
@ -107,7 +107,7 @@ class ContextCachingEndpoints(VertexBase):
Returns
token, url
"""
auth_header: str | None
auth_header: str | dict[str, str | None] | None
if custom_llm_provider == "gemini":
auth_header = {"x-goog-api-key": gemini_api_key}
endpoint = "cachedContents"
@ -129,7 +129,7 @@ class ContextCachingEndpoints(VertexBase):
gemini_api_key=gemini_api_key,
endpoint=endpoint,
stream=None,
auth_header=auth_header,
auth_header=cast(str | None, auth_header),
url=url,
model=model,
vertex_project=vertex_project,
@ -428,21 +428,35 @@ class ContextCachingEndpoints(VertexBase):
client = client
## CHECK IF CACHED ALREADY
generated_cache_key: Final = local_cache_obj.get_cache_key(
messages=cached_messages, tools=tools, tool_choice=tool_choice, model=model
)
google_cache_name: Final = self.check_cache(
cache_key=generated_cache_key,
client=client,
headers=headers,
api_key=api_key,
api_base=api_base,
logging_obj=logging_obj,
custom_llm_provider=custom_llm_provider,
vertex_project=vertex_project,
vertex_location=vertex_location,
vertex_auth_header=vertex_auth_header,
model=model,
cache_key_kwargs: Final[dict[str, object]] = {
"messages": cached_messages,
"tools": tools,
"tool_choice": tool_choice,
"model": model,
}
generated_cache_key: Final = local_cache_obj.get_cache_key(**cache_key_kwargs)
if not generated_cache_key:
if tools is not None:
optional_params["tools"] = tools
if tool_choice is not None:
optional_params["tool_choice"] = tool_choice
return messages, optional_params, None
google_cache_name: Final = (
None
if generated_cache_key is None
else self.check_cache(
cache_key=generated_cache_key,
client=client,
headers=headers,
api_key=api_key,
api_base=api_base,
logging_obj=logging_obj,
custom_llm_provider=custom_llm_provider,
vertex_project=vertex_project,
vertex_location=vertex_location,
vertex_auth_header=vertex_auth_header,
model=model,
)
)
if google_cache_name:
return non_cached_messages, optional_params, google_cache_name
@ -589,9 +603,19 @@ class ContextCachingEndpoints(VertexBase):
client = client
## CHECK IF CACHED ALREADY
generated_cache_key: Final = local_cache_obj.get_cache_key(
messages=cached_messages, tools=tools, tool_choice=tool_choice, model=model
)
cache_key_kwargs: Final[dict[str, object]] = {
"messages": cached_messages,
"tools": tools,
"tool_choice": tool_choice,
"model": model,
}
generated_cache_key: Final = local_cache_obj.get_cache_key(**cache_key_kwargs)
if not generated_cache_key:
if tools is not None:
optional_params["tools"] = tools
if tool_choice is not None:
optional_params["tool_choice"] = tool_choice
return messages, optional_params, None
scope: Final = KnownCachedContents.scope(
url, custom_llm_provider, headers, vertex_project=vertex_project, vertex_location=vertex_location
)

4
litellm/proxy/hooks/batch_redis_get.py
View file

@ -97,7 +97,7 @@ class _PROXY_BatchRedisRequests(CustomLogger):
return None
cache_key: str | None = None
if "cache_key" in kwargs:
cache_key = litellm.cache._prepare_explicit_cache_key( # rebind-ok: normalize explicit caller key
cache_key = litellm.cache.prepare_explicit_cache_key( # rebind-ok: normalize explicit caller key
kwargs["cache_key"], **kwargs
)
else:
@ -105,7 +105,7 @@ class _PROXY_BatchRedisRequests(CustomLogger):
*args, **kwargs
) # returns "<cache_key_name>:<hash>" - we pass redis_namespace in async_pre_call_hook. Done to avoid rewriting the async_set_cache logic
if cache_key is not None and self.in_memory_cache is not None and litellm.cache is not None:
if cache_key is not None and self.in_memory_cache is not None:
cache_control_args: Final = kwargs.get("cache", {})
max_age: Final = cache_control_args.get("s-max-age", cache_control_args.get("s-maxage", float("inf")))
cached_result = self.in_memory_cache.get_cache(cache_key, *args, **kwargs)

2
litellm/proxy/spend_tracking/spend_tracking_utils.py
View file

@ -848,7 +848,7 @@ def get_logging_payload(
elif litellm_params.get("preset_cache_key") is not None:
cache_key = litellm_params["preset_cache_key"]
else:
cache_key = litellm.cache.get_cache_key(**kwargs)
cache_key = litellm.cache.get_cache_key(**cast(dict[str, object], kwargs)) or "Cache OFF"
if cache_hit is True or response_served_from_cache:
import time