mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-14 23:21:35 +00:00
Three process-lifetime retention points kept full request payloads (messages included) alive after the request finished. Under bursts of large-token traffic (~73K tokens/request mean) this presented as stepwise RSS growth that never returned to baseline, ending in OOM: 1. Logging.pre_call/post_call stored their entire locals() (messages, the Logging object, complete_input_dict) in the module-level litellm.error_logs dict, pinning the most recent request's payload per worker forever. Nothing reads that dict; the writes are removed. 2. LLMCachingHandler.request_kwargs kept litellm_logging_obj inside the stored kwargs while the handler itself hangs off logging_obj._llm_caching_handler, closing a reference cycle (Logging -> LLMCachingHandler -> kwargs -> Logging). Cyclic payloads are only reclaimed by generational GC, so megabytes of dead request data lingered until a rare gen-2 pass, and the transient copies fragment the allocator into a permanent RSS high-water mark. The handler now drops litellm_logging_obj from its stored kwargs; the caching layer never reads it. 3. The router stored every request's kwargs in the ITPM/OTPM contextvar even when no deployment configures itpm/otpm. Pooled resources created mid-request (e.g. redis connections) capture the asyncio context, extending that pin far past the request. The slot is now populated only for deployments with io token limits and overwritten with None otherwise. Live-proxy verification (bursts of 30 x ~300KB requests, PII guardrail + prometheus + redis cache): unfixed grows 16-29MB per burst without release; fixed grows under 1MB per burst after warmup and flattens. Resolves LIT-4434
1198 lines
48 KiB
Python
1198 lines
48 KiB
Python
"""
|
|
This contains LLMCachingHandler
|
|
|
|
This exposes two methods:
|
|
- async_get_cache
|
|
- async_set_cache
|
|
|
|
This file is a wrapper around caching.py
|
|
|
|
This class is used to handle caching logic specific for LLM API requests (completion / embedding / text_completion / transcription etc)
|
|
|
|
It utilizes the (RedisCache, s3Cache, RedisSemanticCache, QdrantSemanticCache, InMemoryCache, DiskCache) based on what the user has setup
|
|
|
|
In each method it will call the appropriate method from caching.py
|
|
"""
|
|
|
|
import asyncio
|
|
import datetime
|
|
import inspect
|
|
import time
|
|
from typing import (
|
|
TYPE_CHECKING,
|
|
Any,
|
|
AsyncGenerator,
|
|
Callable,
|
|
Dict,
|
|
Generator,
|
|
List,
|
|
Optional,
|
|
Tuple,
|
|
Union,
|
|
)
|
|
|
|
from pydantic import BaseModel
|
|
|
|
import litellm
|
|
from litellm._logging import print_verbose, verbose_logger
|
|
from litellm.caching import InMemoryCache
|
|
from litellm.caching.caching import S3Cache
|
|
from litellm.litellm_core_utils.llm_response_utils.response_metadata import (
|
|
update_response_metadata,
|
|
)
|
|
from litellm.litellm_core_utils.logging_utils import (
|
|
_assemble_complete_response_from_streaming_chunks,
|
|
)
|
|
from litellm.types.caching import CachedEmbedding
|
|
from litellm.types.llms.openai import ResponsesAPIResponse
|
|
from litellm.types.rerank import RerankResponse
|
|
from litellm.types.utils import (
|
|
CachingDetails,
|
|
CallTypes,
|
|
Embedding,
|
|
EmbeddingResponse,
|
|
ModelResponse,
|
|
TextCompletionResponse,
|
|
TranscriptionResponse,
|
|
Usage,
|
|
)
|
|
|
|
if TYPE_CHECKING:
|
|
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
|
from litellm.types.utils import PromptTokensDetailsWrapper
|
|
else:
|
|
LiteLLMLoggingObj = Any
|
|
|
|
|
|
from litellm.litellm_core_utils.core_helpers import (
|
|
_get_parent_otel_span_from_kwargs,
|
|
)
|
|
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
|
|
|
|
|
|
class CachingHandlerResponse(BaseModel):
|
|
"""
|
|
This is the response object for the caching handler. We need to separate embedding cached responses and (completion / text_completion / transcription) cached responses
|
|
|
|
For embeddings there can be a cache hit for some of the inputs in the list and a cache miss for others
|
|
"""
|
|
|
|
cached_result: Optional[Any] = None
|
|
final_embedding_cached_response: Optional[EmbeddingResponse] = None
|
|
embedding_all_elements_cache_hit: bool = False # this is set to True when all elements in the list have a cache hit in the embedding cache, if true return the final_embedding_cached_response no need to make an API call
|
|
|
|
|
|
in_memory_cache_obj = InMemoryCache()
|
|
|
|
|
|
def _drop_logging_obj_from_kwargs(request_kwargs: dict[str, object]) -> dict[str, object]:
|
|
"""
|
|
The caching handler is stored on the Logging object
|
|
(``logging_obj._llm_caching_handler``), so keeping ``litellm_logging_obj``
|
|
inside ``request_kwargs`` closes a reference cycle
|
|
(Logging -> LLMCachingHandler -> kwargs -> Logging) that keeps the full
|
|
request payload (messages included) alive until a generational GC pass
|
|
instead of being freed by refcount when the request ends. Nothing in the
|
|
caching layer reads the logging object from these kwargs; cache-key
|
|
generation ignores litellm-internal params.
|
|
"""
|
|
if "litellm_logging_obj" not in request_kwargs:
|
|
return request_kwargs
|
|
return {k: v for k, v in request_kwargs.items() if k != "litellm_logging_obj"}
|
|
|
|
|
|
def _is_chat_completion_cached_dict(cached_result: dict) -> bool:
|
|
cached_id = cached_result.get("id")
|
|
if isinstance(cached_id, str) and cached_id.startswith("chatcmpl"):
|
|
return True
|
|
obj = cached_result.get("object")
|
|
if isinstance(obj, str):
|
|
return obj.startswith("chat.completion")
|
|
return "choices" in cached_result
|
|
|
|
|
|
def _should_defer_streaming_cache_hit_callbacks(*, kwargs: Dict[str, Any]) -> bool:
|
|
"""
|
|
When stream=True, do not run success callbacks at cache-hit time.
|
|
|
|
Cached chat/text completion replay uses CustomStreamWrapper; cached Responses
|
|
replay uses CachedResponsesAPIStreamingIterator. Both invoke logging success
|
|
handlers when the stream finishes; firing them here too would double-count
|
|
spend and callback records.
|
|
"""
|
|
return kwargs.get("stream", False) is True
|
|
|
|
|
|
class LLMCachingHandler:
|
|
def __init__(
|
|
self,
|
|
original_function: Callable,
|
|
request_kwargs: Dict[str, Any],
|
|
start_time: datetime.datetime,
|
|
):
|
|
from litellm.caching import DualCache, RedisCache
|
|
|
|
self.async_streaming_chunks: List[ModelResponse] = []
|
|
self.sync_streaming_chunks: List[ModelResponse] = []
|
|
self.request_kwargs = _drop_logging_obj_from_kwargs(request_kwargs)
|
|
self.preset_cache_key: Optional[str] = None
|
|
self.original_function = original_function
|
|
self.start_time = start_time
|
|
if litellm.cache is not None and isinstance(litellm.cache.cache, RedisCache):
|
|
self.dual_cache: Optional[DualCache] = DualCache(
|
|
redis_cache=litellm.cache.cache,
|
|
in_memory_cache=in_memory_cache_obj,
|
|
)
|
|
else:
|
|
self.dual_cache = None
|
|
pass
|
|
|
|
async def _async_get_cache(
|
|
self,
|
|
model: str,
|
|
original_function: Callable,
|
|
logging_obj: LiteLLMLoggingObj,
|
|
start_time: datetime.datetime,
|
|
call_type: str,
|
|
kwargs: Dict[str, Any],
|
|
args: Optional[Tuple[Any, ...]] = None,
|
|
) -> Optional[CachingHandlerResponse]:
|
|
"""
|
|
Internal method to get from the cache.
|
|
Handles different call types (embeddings, chat/completions, text_completion, transcription)
|
|
and accordingly returns the cached response
|
|
|
|
Args:
|
|
model: str:
|
|
original_function: Callable:
|
|
logging_obj: LiteLLMLoggingObj:
|
|
start_time: datetime.datetime:
|
|
call_type: str:
|
|
kwargs: Dict[str, Any]:
|
|
args: Optional[Tuple[Any, ...]] = None:
|
|
|
|
|
|
Returns:
|
|
CachingHandlerResponse:
|
|
Raises:
|
|
None
|
|
"""
|
|
# Check if caching should be performed BEFORE doing expensive operations
|
|
if (
|
|
(kwargs.get("caching", None) is None and litellm.cache is not None) or kwargs.get("caching", False) is True
|
|
) and (
|
|
kwargs.get("cache", {}).get("no-cache", False) is not True
|
|
): # allow users to control returning cached responses from the completion function
|
|
args = args or ()
|
|
final_embedding_cached_response: Optional[EmbeddingResponse] = None
|
|
embedding_all_elements_cache_hit: bool = False
|
|
cached_result: Optional[Any] = None
|
|
kwargs = kwargs.copy()
|
|
#########################################################
|
|
# Init cache timing metrics
|
|
#########################################################
|
|
cache_check_start_time = time.perf_counter()
|
|
cache_check_end_time: Optional[float] = None
|
|
#########################################################
|
|
parent_otel_span = _get_parent_otel_span_from_kwargs(kwargs)
|
|
kwargs["parent_otel_span"] = parent_otel_span
|
|
|
|
if litellm.cache is not None and self._is_call_type_supported_by_cache(original_function=original_function):
|
|
verbose_logger.debug("Checking Async Cache")
|
|
cached_result = await self._retrieve_from_cache(
|
|
call_type=call_type,
|
|
kwargs=kwargs,
|
|
args=args,
|
|
)
|
|
cache_check_end_time = time.perf_counter()
|
|
|
|
if cached_result is not None and not isinstance(cached_result, list):
|
|
verbose_logger.debug("Cache Hit!")
|
|
cache_hit = True
|
|
end_time = datetime.datetime.now()
|
|
model, custom_llm_provider, _, _ = litellm.get_llm_provider(
|
|
model=model,
|
|
custom_llm_provider=kwargs.get("custom_llm_provider", None),
|
|
api_base=kwargs.get("api_base", None),
|
|
api_key=kwargs.get("api_key", None),
|
|
)
|
|
cache_duration_ms = (cache_check_end_time - cache_check_start_time) * 1000
|
|
self._update_litellm_logging_obj_environment(
|
|
logging_obj=logging_obj,
|
|
model=model,
|
|
kwargs=kwargs,
|
|
cached_result=cached_result,
|
|
is_async=True,
|
|
custom_llm_provider=custom_llm_provider,
|
|
cache_duration_ms=cache_duration_ms,
|
|
)
|
|
|
|
call_type = original_function.__name__
|
|
|
|
cached_result = self._convert_cached_result_to_model_response(
|
|
cached_result=cached_result,
|
|
call_type=call_type,
|
|
kwargs=kwargs,
|
|
logging_obj=logging_obj,
|
|
model=model,
|
|
custom_llm_provider=kwargs.get("custom_llm_provider", None),
|
|
args=args,
|
|
)
|
|
if not _should_defer_streaming_cache_hit_callbacks(kwargs=kwargs):
|
|
# LOG SUCCESS
|
|
self._async_log_cache_hit_on_callbacks(
|
|
logging_obj=logging_obj,
|
|
cached_result=cached_result,
|
|
start_time=start_time,
|
|
end_time=end_time,
|
|
cache_hit=cache_hit,
|
|
)
|
|
cache_key = (
|
|
self.preset_cache_key
|
|
or self.request_kwargs.get("cache_key")
|
|
or litellm.cache.get_cache_key(**self.request_kwargs)
|
|
)
|
|
if hasattr(cached_result, "_hidden_params"):
|
|
cached_result._hidden_params["cache_key"] = cache_key # type: ignore
|
|
return CachingHandlerResponse(cached_result=cached_result)
|
|
elif (
|
|
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.
|
|
):
|
|
(
|
|
final_embedding_cached_response,
|
|
embedding_all_elements_cache_hit,
|
|
) = self._process_async_embedding_cached_response(
|
|
final_embedding_cached_response=final_embedding_cached_response,
|
|
cached_result=cached_result,
|
|
kwargs=kwargs,
|
|
logging_obj=logging_obj,
|
|
start_time=start_time,
|
|
model=model,
|
|
)
|
|
return CachingHandlerResponse(
|
|
final_embedding_cached_response=final_embedding_cached_response,
|
|
embedding_all_elements_cache_hit=embedding_all_elements_cache_hit,
|
|
)
|
|
|
|
verbose_logger.debug(f"CACHE RESULT: {cached_result}")
|
|
return CachingHandlerResponse(
|
|
cached_result=cached_result,
|
|
final_embedding_cached_response=final_embedding_cached_response,
|
|
)
|
|
# Caching disabled - return None to indicate no caching attempted
|
|
return None
|
|
|
|
def _sync_get_cache(
|
|
self,
|
|
model: str,
|
|
original_function: Callable,
|
|
logging_obj: LiteLLMLoggingObj,
|
|
start_time: datetime.datetime,
|
|
call_type: str,
|
|
kwargs: Dict[str, Any],
|
|
args: Optional[Tuple[Any, ...]] = None,
|
|
) -> CachingHandlerResponse:
|
|
cached_result: Optional[Any] = None
|
|
|
|
# Check if caching should be performed BEFORE doing expensive kwargs copy
|
|
if litellm.cache is not None and self._is_call_type_supported_by_cache(original_function=original_function):
|
|
args = args or ()
|
|
# Now that we confirmed caching will happen, prepare kwargs
|
|
new_kwargs = kwargs.copy()
|
|
new_kwargs.update(
|
|
convert_args_to_kwargs(
|
|
self.original_function,
|
|
args,
|
|
)
|
|
)
|
|
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)
|
|
self.request_kwargs = _drop_logging_obj_from_kwargs(new_kwargs)
|
|
print_verbose("Checking Sync Cache")
|
|
cached_result = litellm.cache.get_cache(**new_kwargs)
|
|
if cached_result is not None:
|
|
if "detail" in cached_result:
|
|
# implies an error occurred
|
|
pass
|
|
else:
|
|
call_type = original_function.__name__
|
|
cached_result = self._convert_cached_result_to_model_response(
|
|
cached_result=cached_result,
|
|
call_type=call_type,
|
|
kwargs=kwargs,
|
|
logging_obj=logging_obj,
|
|
model=model,
|
|
custom_llm_provider=kwargs.get("custom_llm_provider", None),
|
|
args=args,
|
|
)
|
|
|
|
# LOG SUCCESS
|
|
cache_hit = True
|
|
end_time = datetime.datetime.now()
|
|
(
|
|
model,
|
|
custom_llm_provider,
|
|
dynamic_api_key,
|
|
api_base,
|
|
) = litellm.get_llm_provider(
|
|
model=model or "",
|
|
custom_llm_provider=kwargs.get("custom_llm_provider", None),
|
|
api_base=kwargs.get("api_base", None),
|
|
api_key=kwargs.get("api_key", None),
|
|
)
|
|
self._update_litellm_logging_obj_environment(
|
|
logging_obj=logging_obj,
|
|
model=f"{custom_llm_provider}/{model}",
|
|
kwargs=kwargs,
|
|
cached_result=cached_result,
|
|
is_async=False,
|
|
)
|
|
|
|
if not _should_defer_streaming_cache_hit_callbacks(kwargs=kwargs):
|
|
logging_obj.handle_sync_success_callbacks_for_async_calls(
|
|
result=cached_result,
|
|
start_time=start_time,
|
|
end_time=end_time,
|
|
cache_hit=cache_hit,
|
|
)
|
|
cache_key = (
|
|
self.preset_cache_key
|
|
or self.request_kwargs.get("cache_key")
|
|
or litellm.cache.get_cache_key(**self.request_kwargs)
|
|
)
|
|
if hasattr(cached_result, "_hidden_params"):
|
|
cached_result._hidden_params["cache_key"] = cache_key # type: ignore
|
|
return CachingHandlerResponse(cached_result=cached_result)
|
|
return CachingHandlerResponse(cached_result=cached_result)
|
|
|
|
def handle_kwargs_input_list_or_str(self, kwargs: Dict[str, Any]) -> List[str]:
|
|
"""
|
|
Handles the input of kwargs['input'] being a list or a string
|
|
"""
|
|
if isinstance(kwargs["input"], str):
|
|
return [kwargs["input"]]
|
|
elif isinstance(kwargs["input"], list):
|
|
return kwargs["input"]
|
|
else:
|
|
raise ValueError("input must be a string or a list")
|
|
|
|
def _extract_model_from_cached_results(self, non_null_list: List[Tuple[int, CachedEmbedding]]) -> Optional[str]:
|
|
"""
|
|
Helper method to extract the model name from cached results.
|
|
|
|
Args:
|
|
non_null_list: List of (idx, cr) tuples where cr is the cached result dict
|
|
|
|
Returns:
|
|
Optional[str]: The model name if found, None otherwise
|
|
"""
|
|
for _, cr in non_null_list:
|
|
if isinstance(cr, dict) and cr.get("model"):
|
|
return cr["model"]
|
|
return None
|
|
|
|
def _process_async_embedding_cached_response(
|
|
self,
|
|
final_embedding_cached_response: Optional[EmbeddingResponse],
|
|
cached_result: List[Optional[CachedEmbedding]],
|
|
kwargs: Dict[str, Any],
|
|
logging_obj: LiteLLMLoggingObj,
|
|
start_time: datetime.datetime,
|
|
model: str,
|
|
) -> Tuple[Optional[EmbeddingResponse], bool]:
|
|
"""
|
|
Returns the final embedding cached response and a boolean indicating if all elements in the list have a cache hit
|
|
|
|
For embedding responses, there can be a cache hit for some of the inputs in the list and a cache miss for others
|
|
This function processes the cached embedding responses and returns the final embedding cached response and a boolean indicating if all elements in the list have a cache hit
|
|
|
|
Args:
|
|
final_embedding_cached_response: Optional[EmbeddingResponse]:
|
|
cached_result: List[Optional[Dict[str, Any]]]:
|
|
kwargs: Dict[str, Any]:
|
|
logging_obj: LiteLLMLoggingObj:
|
|
start_time: datetime.datetime:
|
|
model: str:
|
|
|
|
Returns:
|
|
Tuple[Optional[EmbeddingResponse], bool]:
|
|
Returns the final embedding cached response and a boolean indicating if all elements in the list have a cache hit
|
|
|
|
|
|
"""
|
|
embedding_all_elements_cache_hit: bool = False
|
|
remaining_list = []
|
|
non_null_list = []
|
|
kwargs_input_as_list = self.handle_kwargs_input_list_or_str(kwargs)
|
|
for idx, cr in enumerate(cached_result):
|
|
if cr is None:
|
|
remaining_list.append(kwargs_input_as_list[idx])
|
|
else:
|
|
non_null_list.append((idx, cr))
|
|
kwargs["input"] = remaining_list
|
|
if len(non_null_list) > 0:
|
|
# Use the model from the first non-null cached result, fallback to kwargs if not present
|
|
model_name = self._extract_model_from_cached_results(non_null_list)
|
|
if not model_name:
|
|
model_name = kwargs.get("model")
|
|
final_embedding_cached_response = EmbeddingResponse(
|
|
model=model_name,
|
|
data=[None] * len(kwargs_input_as_list),
|
|
)
|
|
final_embedding_cached_response._hidden_params["cache_hit"] = True
|
|
|
|
prompt_tokens = 0
|
|
aggregated_details: Optional[dict] = None
|
|
for val in non_null_list:
|
|
idx, cr = val # (idx, cr) tuple
|
|
if cr is not None:
|
|
embedding_data = cr.get("embedding")
|
|
if embedding_data is not None:
|
|
final_embedding_cached_response.data[idx] = Embedding(
|
|
embedding=embedding_data,
|
|
index=idx,
|
|
object="embedding",
|
|
)
|
|
cached_prompt_tokens = cr.get("prompt_tokens")
|
|
if cached_prompt_tokens is not None:
|
|
prompt_tokens += cached_prompt_tokens
|
|
elif isinstance(kwargs_input_as_list[idx], str):
|
|
from litellm.utils import token_counter
|
|
|
|
prompt_tokens += token_counter(text=kwargs_input_as_list[idx], count_response_tokens=True)
|
|
# Aggregate prompt_tokens_details from cached items
|
|
item_details = cr.get("prompt_tokens_details")
|
|
if item_details:
|
|
if aggregated_details is None:
|
|
aggregated_details = {}
|
|
for key, value in item_details.items():
|
|
if isinstance(value, (int, float)):
|
|
aggregated_details[key] = aggregated_details.get(key, 0) + value
|
|
else:
|
|
aggregated_details[key] = value
|
|
|
|
## USAGE
|
|
prompt_tokens_details: Optional["PromptTokensDetailsWrapper"] = None
|
|
if aggregated_details:
|
|
from litellm.types.utils import PromptTokensDetailsWrapper
|
|
|
|
try:
|
|
prompt_tokens_details = PromptTokensDetailsWrapper(**aggregated_details)
|
|
except Exception:
|
|
prompt_tokens_details = None
|
|
usage = Usage(
|
|
prompt_tokens=prompt_tokens,
|
|
completion_tokens=0,
|
|
total_tokens=prompt_tokens,
|
|
prompt_tokens_details=prompt_tokens_details,
|
|
)
|
|
final_embedding_cached_response.usage = usage
|
|
if len(remaining_list) == 0:
|
|
# LOG SUCCESS
|
|
cache_hit = True
|
|
embedding_all_elements_cache_hit = True
|
|
end_time = datetime.datetime.now()
|
|
(
|
|
model,
|
|
custom_llm_provider,
|
|
dynamic_api_key,
|
|
api_base,
|
|
) = litellm.get_llm_provider(
|
|
model=model,
|
|
custom_llm_provider=kwargs.get("custom_llm_provider", None),
|
|
api_base=kwargs.get("api_base", None),
|
|
api_key=kwargs.get("api_key", None),
|
|
)
|
|
|
|
self._update_litellm_logging_obj_environment(
|
|
logging_obj=logging_obj,
|
|
model=model,
|
|
kwargs=kwargs,
|
|
cached_result=final_embedding_cached_response,
|
|
is_async=True,
|
|
is_embedding=True,
|
|
)
|
|
self._async_log_cache_hit_on_callbacks(
|
|
logging_obj=logging_obj,
|
|
cached_result=final_embedding_cached_response,
|
|
start_time=start_time,
|
|
end_time=end_time,
|
|
cache_hit=cache_hit,
|
|
)
|
|
return final_embedding_cached_response, embedding_all_elements_cache_hit
|
|
return final_embedding_cached_response, embedding_all_elements_cache_hit
|
|
|
|
def combine_usage(self, usage1: Usage, usage2: Usage) -> Usage:
|
|
return Usage(
|
|
prompt_tokens=usage1.prompt_tokens + usage2.prompt_tokens,
|
|
completion_tokens=usage1.completion_tokens + usage2.completion_tokens,
|
|
total_tokens=usage1.total_tokens + usage2.total_tokens,
|
|
prompt_tokens_details=self._merge_prompt_tokens_details(
|
|
usage1.prompt_tokens_details,
|
|
usage2.prompt_tokens_details,
|
|
),
|
|
)
|
|
|
|
def _merge_prompt_tokens_details(
|
|
self,
|
|
details1: Optional["PromptTokensDetailsWrapper"],
|
|
details2: Optional["PromptTokensDetailsWrapper"],
|
|
) -> Optional["PromptTokensDetailsWrapper"]:
|
|
"""Merge two PromptTokensDetailsWrapper objects by summing numeric fields."""
|
|
if details1 is None and details2 is None:
|
|
return None
|
|
if details1 is None:
|
|
return details2
|
|
if details2 is None:
|
|
return details1
|
|
|
|
dict1 = details1.model_dump(exclude_none=True) if hasattr(details1, "model_dump") else {}
|
|
dict2 = details2.model_dump(exclude_none=True) if hasattr(details2, "model_dump") else {}
|
|
|
|
merged: dict = {}
|
|
for key in set(dict1.keys()) | set(dict2.keys()):
|
|
v1 = dict1.get(key, 0)
|
|
v2 = dict2.get(key, 0)
|
|
if isinstance(v1, (int, float)) and isinstance(v2, (int, float)):
|
|
merged[key] = v1 + v2
|
|
elif isinstance(v1, dict) and isinstance(v2, dict):
|
|
# Recursively merge nested dicts (e.g. cache_creation_token_details)
|
|
nested: dict = {}
|
|
for nk in set(v1.keys()) | set(v2.keys()):
|
|
nv1 = v1.get(nk, 0)
|
|
nv2 = v2.get(nk, 0)
|
|
if isinstance(nv1, (int, float)) and isinstance(nv2, (int, float)):
|
|
nested[nk] = nv1 + nv2
|
|
elif nv1:
|
|
nested[nk] = nv1
|
|
else:
|
|
nested[nk] = nv2
|
|
merged[key] = nested
|
|
elif v1:
|
|
merged[key] = v1
|
|
else:
|
|
merged[key] = v2
|
|
|
|
if not merged:
|
|
return None
|
|
|
|
from litellm.types.utils import PromptTokensDetailsWrapper
|
|
|
|
try:
|
|
return PromptTokensDetailsWrapper(**merged)
|
|
except Exception:
|
|
return None
|
|
|
|
def _combine_cached_embedding_response_with_api_result(
|
|
self,
|
|
_caching_handler_response: CachingHandlerResponse,
|
|
embedding_response: EmbeddingResponse,
|
|
start_time: datetime.datetime,
|
|
end_time: datetime.datetime,
|
|
) -> EmbeddingResponse:
|
|
"""
|
|
Combines the cached embedding response with the API EmbeddingResponse
|
|
|
|
For caching there can be a cache hit for some of the inputs in the list and a cache miss for others
|
|
This function combines the cached embedding response with the API EmbeddingResponse
|
|
|
|
Args:
|
|
caching_handler_response: CachingHandlerResponse:
|
|
embedding_response: EmbeddingResponse:
|
|
|
|
Returns:
|
|
EmbeddingResponse:
|
|
"""
|
|
if _caching_handler_response.final_embedding_cached_response is None:
|
|
return embedding_response
|
|
|
|
idx = 0
|
|
final_data_list = []
|
|
for item in _caching_handler_response.final_embedding_cached_response.data:
|
|
if item is None and embedding_response.data is not None:
|
|
final_data_list.append(embedding_response.data[idx])
|
|
idx += 1
|
|
else:
|
|
final_data_list.append(item)
|
|
|
|
_caching_handler_response.final_embedding_cached_response.data = final_data_list
|
|
_caching_handler_response.final_embedding_cached_response._hidden_params["cache_hit"] = True
|
|
_caching_handler_response.final_embedding_cached_response._response_ms = (
|
|
end_time - start_time
|
|
).total_seconds() * 1000
|
|
|
|
## USAGE
|
|
if (
|
|
_caching_handler_response.final_embedding_cached_response.usage is not None
|
|
and embedding_response.usage is not None
|
|
):
|
|
_caching_handler_response.final_embedding_cached_response.usage = self.combine_usage(
|
|
usage1=_caching_handler_response.final_embedding_cached_response.usage,
|
|
usage2=embedding_response.usage,
|
|
)
|
|
|
|
return _caching_handler_response.final_embedding_cached_response
|
|
|
|
def _async_log_cache_hit_on_callbacks(
|
|
self,
|
|
logging_obj: LiteLLMLoggingObj,
|
|
cached_result: Any,
|
|
start_time: datetime.datetime,
|
|
end_time: datetime.datetime,
|
|
cache_hit: bool,
|
|
):
|
|
"""
|
|
Helper function to log the success of a cached result on callbacks
|
|
|
|
Args:
|
|
logging_obj (LiteLLMLoggingObj): The logging object.
|
|
cached_result: The cached result.
|
|
start_time (datetime): The start time of the operation.
|
|
end_time (datetime): The end time of the operation.
|
|
cache_hit (bool): Whether it was a cache hit.
|
|
"""
|
|
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
|
|
|
|
GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(
|
|
async_coroutine=logging_obj.async_success_handler(
|
|
result=cached_result,
|
|
start_time=start_time,
|
|
end_time=end_time,
|
|
cache_hit=cache_hit,
|
|
)
|
|
)
|
|
|
|
logging_obj.handle_sync_success_callbacks_for_async_calls(
|
|
result=cached_result,
|
|
start_time=start_time,
|
|
end_time=end_time,
|
|
cache_hit=cache_hit,
|
|
)
|
|
|
|
async def _retrieve_from_cache(
|
|
self, call_type: str, kwargs: Dict[str, Any], args: Tuple[Any, ...]
|
|
) -> Optional[Any]:
|
|
"""
|
|
Internal method to
|
|
- get cache key
|
|
- check what type of cache is used - Redis, RedisSemantic, Qdrant, S3
|
|
- async get cache value
|
|
- return the cached value
|
|
|
|
Args:
|
|
call_type: str:
|
|
kwargs: Dict[str, Any]:
|
|
args: Optional[Tuple[Any, ...]] = None:
|
|
|
|
Returns:
|
|
Optional[Any]:
|
|
Raises:
|
|
None
|
|
"""
|
|
if litellm.cache is None:
|
|
return None
|
|
|
|
new_kwargs = kwargs.copy()
|
|
new_kwargs.update(
|
|
convert_args_to_kwargs(
|
|
self.original_function,
|
|
args,
|
|
)
|
|
)
|
|
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)
|
|
self.request_kwargs = _drop_logging_obj_from_kwargs(new_kwargs)
|
|
cached_result: Optional[Any] = None
|
|
if call_type == CallTypes.aembedding.value:
|
|
if isinstance(new_kwargs["input"], str):
|
|
new_kwargs["input"] = [new_kwargs["input"]]
|
|
elif not isinstance(new_kwargs["input"], list):
|
|
raise ValueError("input must be a string or a list")
|
|
tasks = []
|
|
for idx, i in enumerate(new_kwargs["input"]):
|
|
preset_cache_key = litellm.cache.get_cache_key(**{**new_kwargs, "input": i})
|
|
tasks.append(
|
|
litellm.cache.async_get_cache(
|
|
cache_key=preset_cache_key,
|
|
dynamic_cache_object=self.dual_cache,
|
|
)
|
|
)
|
|
cached_result = await asyncio.gather(*tasks)
|
|
## check if cached result is None ##
|
|
if cached_result is not None and 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
|
|
else:
|
|
request_kwargs = new_kwargs.copy()
|
|
request_cache_key = request_kwargs.pop("cache_key", 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)
|
|
cached_result = await litellm.cache.async_get_cache(
|
|
dynamic_cache_object=self.dual_cache,
|
|
cache_key=self.preset_cache_key,
|
|
**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)
|
|
cached_result = litellm.cache.get_cache(
|
|
dynamic_cache_object=self.dual_cache,
|
|
cache_key=self.preset_cache_key,
|
|
**request_kwargs,
|
|
)
|
|
return cached_result
|
|
|
|
def _convert_cached_result_to_model_response(
|
|
self,
|
|
cached_result: Any,
|
|
call_type: str,
|
|
kwargs: Dict[str, Any],
|
|
logging_obj: LiteLLMLoggingObj,
|
|
model: str,
|
|
args: Tuple[Any, ...],
|
|
custom_llm_provider: Optional[str] = None,
|
|
) -> Optional[
|
|
Union[
|
|
ModelResponse,
|
|
TextCompletionResponse,
|
|
EmbeddingResponse,
|
|
RerankResponse,
|
|
TranscriptionResponse,
|
|
CustomStreamWrapper,
|
|
]
|
|
]:
|
|
"""
|
|
Internal method to process the cached result
|
|
|
|
Checks the call type and converts the cached result to the appropriate model response object
|
|
example if call type is text_completion -> returns TextCompletionResponse object
|
|
|
|
Args:
|
|
cached_result: Any:
|
|
call_type: str:
|
|
kwargs: Dict[str, Any]:
|
|
logging_obj: LiteLLMLoggingObj:
|
|
model: str:
|
|
custom_llm_provider: Optional[str] = None:
|
|
args: Optional[Tuple[Any, ...]] = None:
|
|
|
|
Returns:
|
|
Optional[Any]:
|
|
"""
|
|
from litellm.utils import convert_to_model_response_object
|
|
|
|
if (call_type == CallTypes.acompletion.value or call_type == CallTypes.completion.value) and isinstance(
|
|
cached_result, dict
|
|
):
|
|
if kwargs.get("stream", False) is True:
|
|
cached_result = self._convert_cached_stream_response(
|
|
cached_result=cached_result,
|
|
call_type=call_type,
|
|
logging_obj=logging_obj,
|
|
model=model,
|
|
)
|
|
else:
|
|
cached_result = convert_to_model_response_object(
|
|
response_object=cached_result,
|
|
model_response_object=ModelResponse(),
|
|
)
|
|
if (
|
|
call_type == CallTypes.atext_completion.value or call_type == CallTypes.text_completion.value
|
|
) and isinstance(cached_result, dict):
|
|
if kwargs.get("stream", False) is True:
|
|
cached_result = self._convert_cached_stream_response(
|
|
cached_result=cached_result,
|
|
call_type=call_type,
|
|
logging_obj=logging_obj,
|
|
model=model,
|
|
)
|
|
else:
|
|
cached_result = TextCompletionResponse(**cached_result)
|
|
elif (call_type == CallTypes.aembedding.value or call_type == CallTypes.embedding.value) and isinstance(
|
|
cached_result, dict
|
|
):
|
|
cached_result = convert_to_model_response_object(
|
|
response_object=cached_result,
|
|
model_response_object=EmbeddingResponse(),
|
|
response_type="embedding",
|
|
)
|
|
|
|
elif (call_type == CallTypes.arerank.value or call_type == CallTypes.rerank.value) and isinstance(
|
|
cached_result, dict
|
|
):
|
|
cached_result = convert_to_model_response_object(
|
|
response_object=cached_result,
|
|
model_response_object=None,
|
|
response_type="rerank",
|
|
)
|
|
elif (call_type == CallTypes.atranscription.value or call_type == CallTypes.transcription.value) and isinstance(
|
|
cached_result, dict
|
|
):
|
|
hidden_params = {
|
|
"model": "whisper-1",
|
|
"custom_llm_provider": custom_llm_provider,
|
|
"cache_hit": True,
|
|
}
|
|
cached_result = convert_to_model_response_object(
|
|
response_object=cached_result,
|
|
model_response_object=TranscriptionResponse(),
|
|
response_type="audio_transcription",
|
|
hidden_params=hidden_params,
|
|
)
|
|
elif (call_type == "aresponses" or call_type == "responses") and isinstance(cached_result, dict):
|
|
use_chat_completion_cache = _is_chat_completion_cached_dict(cached_result)
|
|
if use_chat_completion_cache:
|
|
if kwargs.get("stream", False) is True:
|
|
bridge_call_type = (
|
|
CallTypes.acompletion.value if call_type == "aresponses" else CallTypes.completion.value
|
|
)
|
|
cached_result = self._convert_cached_stream_response(
|
|
cached_result=cached_result,
|
|
call_type=bridge_call_type,
|
|
logging_obj=logging_obj,
|
|
model=model,
|
|
)
|
|
else:
|
|
cached_result = convert_to_model_response_object(
|
|
response_object=cached_result,
|
|
model_response_object=ModelResponse(),
|
|
)
|
|
else:
|
|
from litellm.responses.streaming_iterator import (
|
|
CachedResponsesAPIStreamingIterator,
|
|
)
|
|
|
|
response_obj = ResponsesAPIResponse(**cached_result)
|
|
if (
|
|
hasattr(response_obj, "_hidden_params")
|
|
and response_obj._hidden_params is not None
|
|
and isinstance(response_obj._hidden_params, dict)
|
|
):
|
|
response_obj._hidden_params["cache_hit"] = True
|
|
|
|
if kwargs.get("stream", False) is True:
|
|
cached_result = CachedResponsesAPIStreamingIterator(
|
|
response=response_obj,
|
|
logging_obj=logging_obj,
|
|
request_data=kwargs,
|
|
call_type=call_type,
|
|
)
|
|
else:
|
|
cached_result = response_obj
|
|
|
|
if (
|
|
hasattr(cached_result, "_hidden_params")
|
|
and cached_result._hidden_params is not None
|
|
and isinstance(cached_result._hidden_params, dict)
|
|
):
|
|
cached_result._hidden_params["cache_hit"] = True
|
|
|
|
#########################################################
|
|
# Add final timing metrics to the cached result
|
|
#########################################################
|
|
update_response_metadata(
|
|
result=cached_result,
|
|
logging_obj=logging_obj,
|
|
model=model,
|
|
kwargs=kwargs,
|
|
start_time=self.start_time,
|
|
end_time=datetime.datetime.now(),
|
|
)
|
|
return cached_result
|
|
|
|
def _convert_cached_stream_response(
|
|
self,
|
|
cached_result: Any,
|
|
call_type: str,
|
|
logging_obj: LiteLLMLoggingObj,
|
|
model: str,
|
|
) -> CustomStreamWrapper:
|
|
from litellm.utils import (
|
|
CustomStreamWrapper,
|
|
convert_to_streaming_response,
|
|
convert_to_streaming_response_async,
|
|
)
|
|
|
|
_stream_cached_result: Union[AsyncGenerator, Generator]
|
|
if call_type == CallTypes.acompletion.value or call_type == CallTypes.atext_completion.value:
|
|
_stream_cached_result = convert_to_streaming_response_async(
|
|
response_object=cached_result,
|
|
)
|
|
else:
|
|
_stream_cached_result = convert_to_streaming_response(
|
|
response_object=cached_result,
|
|
)
|
|
return CustomStreamWrapper(
|
|
completion_stream=_stream_cached_result,
|
|
model=model,
|
|
custom_llm_provider="cached_response",
|
|
logging_obj=logging_obj,
|
|
)
|
|
|
|
async def async_set_cache(
|
|
self,
|
|
result: Any,
|
|
original_function: Callable,
|
|
kwargs: Dict[str, Any],
|
|
args: Optional[Tuple[Any, ...]] = None,
|
|
):
|
|
"""
|
|
Internal method to check the type of the result & cache used and adds the result to the cache accordingly
|
|
|
|
Args:
|
|
result: Any:
|
|
original_function: Callable:
|
|
kwargs: Dict[str, Any]:
|
|
args: Optional[Tuple[Any, ...]] = None:
|
|
|
|
Returns:
|
|
None
|
|
Raises:
|
|
None
|
|
"""
|
|
from litellm.litellm_core_utils.core_helpers import (
|
|
_get_parent_otel_span_from_kwargs,
|
|
)
|
|
|
|
if litellm.cache is None:
|
|
return
|
|
|
|
new_kwargs = kwargs.copy()
|
|
new_kwargs.update(
|
|
convert_args_to_kwargs(
|
|
original_function,
|
|
args,
|
|
)
|
|
)
|
|
parent_otel_span = _get_parent_otel_span_from_kwargs(new_kwargs)
|
|
new_kwargs["parent_otel_span"] = parent_otel_span
|
|
# [OPTIONAL] ADD TO CACHE
|
|
if self._should_store_result_in_cache(original_function=original_function, kwargs=new_kwargs):
|
|
if (
|
|
isinstance(result, litellm.ModelResponse)
|
|
or isinstance(result, litellm.EmbeddingResponse)
|
|
or isinstance(result, TranscriptionResponse)
|
|
or isinstance(result, RerankResponse)
|
|
or isinstance(result, ResponsesAPIResponse)
|
|
):
|
|
if (
|
|
isinstance(result, EmbeddingResponse)
|
|
and litellm.cache is not None
|
|
and not isinstance(litellm.cache.cache, S3Cache) # s3 doesn't support bulk writing. Exclude.
|
|
):
|
|
asyncio.create_task(
|
|
litellm.cache.async_add_cache_pipeline(
|
|
result, dynamic_cache_object=self.dual_cache, **new_kwargs
|
|
)
|
|
)
|
|
else:
|
|
asyncio.create_task(
|
|
litellm.cache.async_add_cache(
|
|
result.model_dump_json(),
|
|
dynamic_cache_object=self.dual_cache,
|
|
**new_kwargs,
|
|
)
|
|
)
|
|
else:
|
|
asyncio.create_task(litellm.cache.async_add_cache(result, **new_kwargs))
|
|
|
|
def sync_set_cache(
|
|
self,
|
|
result: Any,
|
|
kwargs: Dict[str, Any],
|
|
args: Optional[Tuple[Any, ...]] = None,
|
|
):
|
|
"""
|
|
Sync internal method to add the result to the cache
|
|
"""
|
|
|
|
new_kwargs = kwargs.copy()
|
|
new_kwargs.update(
|
|
convert_args_to_kwargs(
|
|
self.original_function,
|
|
args,
|
|
)
|
|
)
|
|
if litellm.cache is None:
|
|
return
|
|
|
|
if self._should_store_result_in_cache(original_function=self.original_function, kwargs=new_kwargs):
|
|
litellm.cache.add_cache(result, **new_kwargs)
|
|
|
|
return
|
|
|
|
def _should_store_result_in_cache(self, original_function: Callable, kwargs: Dict[str, Any]) -> bool:
|
|
"""
|
|
Helper function to determine if the result should be stored in the cache.
|
|
|
|
Returns:
|
|
bool: True if the result should be stored in the cache, False otherwise.
|
|
"""
|
|
return (
|
|
(litellm.cache is not None)
|
|
and litellm.cache.supported_call_types is not None
|
|
and (str(original_function.__name__) in litellm.cache.supported_call_types)
|
|
and (kwargs.get("cache", {}).get("no-store", False) is not True)
|
|
)
|
|
|
|
def _is_call_type_supported_by_cache(
|
|
self,
|
|
original_function: Callable,
|
|
) -> bool:
|
|
"""
|
|
Helper function to determine if the call type is supported by the cache.
|
|
|
|
call types are acompletion, aembedding, atext_completion, atranscription, arerank
|
|
|
|
Defined on `litellm.types.utils.CallTypes`
|
|
|
|
Returns:
|
|
bool: True if the call type is supported by the cache, False otherwise.
|
|
"""
|
|
if (
|
|
litellm.cache is not None
|
|
and litellm.cache.supported_call_types is not None
|
|
and str(original_function.__name__) in litellm.cache.supported_call_types
|
|
):
|
|
return True
|
|
return False
|
|
|
|
async def _add_streaming_response_to_cache(self, processed_chunk: ModelResponse):
|
|
"""
|
|
Internal method to add the streaming response to the cache
|
|
|
|
|
|
- If 'streaming_chunk' has a 'finish_reason' then assemble a litellm.ModelResponse object
|
|
- Else append the chunk to self.async_streaming_chunks
|
|
|
|
"""
|
|
|
|
complete_streaming_response: Optional[Union[ModelResponse, TextCompletionResponse]] = (
|
|
_assemble_complete_response_from_streaming_chunks(
|
|
result=processed_chunk,
|
|
start_time=self.start_time,
|
|
end_time=datetime.datetime.now(),
|
|
request_kwargs=self.request_kwargs,
|
|
streaming_chunks=self.async_streaming_chunks,
|
|
is_async=True,
|
|
)
|
|
)
|
|
# if a complete_streaming_response is assembled, add it to the cache
|
|
if complete_streaming_response is not None:
|
|
await self.async_set_cache(
|
|
result=complete_streaming_response,
|
|
original_function=self.original_function,
|
|
kwargs=self.request_kwargs,
|
|
)
|
|
|
|
def _sync_add_streaming_response_to_cache(self, processed_chunk: ModelResponse):
|
|
"""
|
|
Sync internal method to add the streaming response to the cache
|
|
"""
|
|
complete_streaming_response: Optional[Union[ModelResponse, TextCompletionResponse]] = (
|
|
_assemble_complete_response_from_streaming_chunks(
|
|
result=processed_chunk,
|
|
start_time=self.start_time,
|
|
end_time=datetime.datetime.now(),
|
|
request_kwargs=self.request_kwargs,
|
|
streaming_chunks=self.sync_streaming_chunks,
|
|
is_async=False,
|
|
)
|
|
)
|
|
|
|
# if a complete_streaming_response is assembled, add it to the cache
|
|
if complete_streaming_response is not None:
|
|
self.sync_set_cache(
|
|
result=complete_streaming_response,
|
|
kwargs=self.request_kwargs,
|
|
)
|
|
|
|
def _update_litellm_logging_obj_environment(
|
|
self,
|
|
logging_obj: LiteLLMLoggingObj,
|
|
model: str,
|
|
kwargs: Dict[str, Any],
|
|
cached_result: Any,
|
|
is_async: bool,
|
|
is_embedding: bool = False,
|
|
custom_llm_provider: Optional[str] = None,
|
|
cache_duration_ms: Optional[float] = None,
|
|
):
|
|
"""
|
|
Helper function to update the LiteLLMLoggingObj environment variables.
|
|
|
|
Args:
|
|
logging_obj (LiteLLMLoggingObj): The logging object to update.
|
|
model (str): The model being used.
|
|
kwargs (Dict[str, Any]): The keyword arguments from the original function call.
|
|
cached_result (Any): The cached result to log.
|
|
is_async (bool): Whether the call is asynchronous or not.
|
|
is_embedding (bool): Whether the call is for embeddings or not.
|
|
custom_llm_provider (Optional[str]): The custom llm provider being used.
|
|
|
|
Returns:
|
|
None
|
|
"""
|
|
litellm_params = {
|
|
"logger_fn": kwargs.get("logger_fn", None),
|
|
"acompletion": is_async,
|
|
"api_base": kwargs.get("api_base", ""),
|
|
"metadata": kwargs.get("metadata", {}),
|
|
"model_info": kwargs.get("model_info", {}),
|
|
"proxy_server_request": kwargs.get("proxy_server_request", None),
|
|
"stream_response": kwargs.get("stream_response", {}),
|
|
"custom_llm_provider": custom_llm_provider,
|
|
}
|
|
|
|
if litellm.cache is not None:
|
|
litellm_params["preset_cache_key"] = litellm.cache._get_preset_cache_key_from_kwargs(**kwargs)
|
|
else:
|
|
litellm_params["preset_cache_key"] = None
|
|
|
|
logging_obj.update_environment_variables(
|
|
model=model,
|
|
user=kwargs.get("user", None),
|
|
optional_params={},
|
|
litellm_params=litellm_params,
|
|
input=(kwargs.get("messages", "") if not is_embedding else kwargs.get("input", "")),
|
|
api_key=kwargs.get("api_key", None),
|
|
original_response=str(cached_result),
|
|
additional_args=None,
|
|
stream=kwargs.get("stream", False),
|
|
custom_llm_provider=custom_llm_provider,
|
|
)
|
|
|
|
logging_obj.caching_details = CachingDetails(
|
|
cache_hit=True,
|
|
cache_duration_ms=cache_duration_ms,
|
|
)
|
|
|
|
|
|
def convert_args_to_kwargs(
|
|
original_function: Callable,
|
|
args: Optional[Tuple[Any, ...]] = None,
|
|
) -> Dict[str, Any]:
|
|
# Get the signature of the original function
|
|
signature = inspect.signature(original_function)
|
|
|
|
# Get parameter names in the order they appear in the original function
|
|
param_names = list(signature.parameters.keys())
|
|
|
|
# Create a mapping of positional arguments to parameter names
|
|
args_to_kwargs = {}
|
|
if args:
|
|
for index, arg in enumerate(args):
|
|
if index < len(param_names):
|
|
param_name = param_names[index]
|
|
args_to_kwargs[param_name] = arg
|
|
|
|
return args_to_kwargs
|