diff --git a/litellm/caching/caching.py b/litellm/caching/caching.py index c16993625173..088e2d03fefc 100644 --- a/litellm/caching/caching.py +++ b/litellm/caching/caching.py @@ -25,8 +25,9 @@ import litellm from litellm._logging import verbose_logger from litellm.litellm_core_utils.core_helpers import _get_parent_otel_span_from_kwargs +from litellm.types.caching import * from litellm.types.services import ServiceLoggerPayload, ServiceTypes -from litellm.types.utils import CachingSupportedCallTypes, all_litellm_params +from litellm.types.utils import all_litellm_params def print_verbose(print_statement): @@ -2125,9 +2126,7 @@ async def async_delete_cache(self, key: str): class Cache: def __init__( self, - type: Optional[ - Literal["local", "redis", "redis-semantic", "s3", "disk", "qdrant-semantic"] - ] = "local", + type: Optional[LiteLLMCacheType] = LiteLLMCacheType.LOCAL, mode: Optional[ CacheMode ] = CacheMode.default_on, # when default_on cache is always on, when default_off cache is opt in @@ -2221,7 +2220,7 @@ def __init__( Returns: None. Cache is set as a litellm param """ - if type == "redis": + if type == LiteLLMCacheType.REDIS: self.cache: BaseCache = RedisCache( host=host, port=port, @@ -2230,7 +2229,7 @@ def __init__( startup_nodes=redis_startup_nodes, **kwargs, ) - elif type == "redis-semantic": + elif type == LiteLLMCacheType.REDIS_SEMANTIC: self.cache = RedisSemanticCache( host=host, port=port, @@ -2240,7 +2239,7 @@ def __init__( embedding_model=redis_semantic_cache_embedding_model, **kwargs, ) - elif type == "qdrant-semantic": + elif type == LiteLLMCacheType.QDRANT_SEMANTIC: self.cache = QdrantSemanticCache( qdrant_api_base=qdrant_api_base, qdrant_api_key=qdrant_api_key, @@ -2249,9 +2248,9 @@ def __init__( quantization_config=qdrant_quantization_config, embedding_model=qdrant_semantic_cache_embedding_model, ) - elif type == "local": + elif type == LiteLLMCacheType.LOCAL: self.cache = InMemoryCache() - elif type == "s3": + elif type == LiteLLMCacheType.S3: self.cache = S3Cache( s3_bucket_name=s3_bucket_name, s3_region_name=s3_region_name, @@ -2266,7 +2265,7 @@ def __init__( s3_path=s3_path, **kwargs, ) - elif type == "disk": + elif type == LiteLLMCacheType.DISK: self.cache = DiskCache(disk_cache_dir=disk_cache_dir) if "cache" not in litellm.input_callback: litellm.input_callback.append("cache") @@ -2281,11 +2280,12 @@ def __init__( self.ttl = ttl self.mode: CacheMode = mode or CacheMode.default_on - if self.type == "local" and default_in_memory_ttl is not None: + if self.type == LiteLLMCacheType.LOCAL and default_in_memory_ttl is not None: self.ttl = default_in_memory_ttl if ( - self.type == "redis" or self.type == "redis-semantic" + self.type == LiteLLMCacheType.REDIS + or self.type == LiteLLMCacheType.REDIS_SEMANTIC ) and default_in_redis_ttl is not None: self.ttl = default_in_redis_ttl @@ -2694,6 +2694,17 @@ async def disconnect(self): if hasattr(self.cache, "disconnect"): await self.cache.disconnect() + def _supports_async(self) -> bool: + """ + Internal method to check if the cache type supports async get/set operations + + Only S3 Cache Does NOT support async operations + + """ + if self.type and self.type == LiteLLMCacheType.S3: + return False + return True + class DiskCache(BaseCache): def __init__(self, disk_cache_dir: Optional[str] = None): @@ -2774,7 +2785,7 @@ def delete_cache(self, key): def enable_cache( - type: Optional[Literal["local", "redis", "s3", "disk"]] = "local", + type: Optional[LiteLLMCacheType] = LiteLLMCacheType.LOCAL, host: Optional[str] = None, port: Optional[str] = None, password: Optional[str] = None, @@ -2832,7 +2843,7 @@ def enable_cache( def update_cache( - type: Optional[Literal["local", "redis", "s3", "disk"]] = "local", + type: Optional[LiteLLMCacheType] = LiteLLMCacheType.LOCAL, host: Optional[str] = None, port: Optional[str] = None, password: Optional[str] = None, diff --git a/litellm/caching/caching_handler.py b/litellm/caching/caching_handler.py index 58de0b42fe26..264fb405baa9 100644 --- a/litellm/caching/caching_handler.py +++ b/litellm/caching/caching_handler.py @@ -13,12 +13,12 @@ import asyncio import datetime import threading -from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple +from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Union from pydantic import BaseModel import litellm -from litellm._logging import print_verbose +from litellm._logging import print_verbose, verbose_logger from litellm.caching.caching import ( Cache, QdrantSemanticCache, @@ -57,7 +57,16 @@ class CachingHandlerResponse(BaseModel): class LLMCachingHandler: - def __init__(self): + def __init__( + self, + original_function: Callable, + request_kwargs: Dict[str, Any], + start_time: datetime.datetime, + ): + self.async_streaming_chunks: List[ModelResponse] = [] + self.request_kwargs = request_kwargs + self.original_function = original_function + self.start_time = start_time pass async def _async_get_cache( @@ -90,11 +99,7 @@ async def _async_get_cache( Raises: None """ - from litellm.utils import ( - CustomStreamWrapper, - convert_to_model_response_object, - convert_to_streaming_response_async, - ) + from litellm.utils import CustomStreamWrapper args = args or () @@ -115,43 +120,13 @@ async def _async_get_cache( in litellm.cache.supported_call_types ): print_verbose("Checking Cache") - if call_type == CallTypes.aembedding.value and isinstance( - kwargs["input"], list - ): - tasks = [] - for idx, i in enumerate(kwargs["input"]): - preset_cache_key = litellm.cache.get_cache_key( - *args, **{**kwargs, "input": i} - ) - tasks.append( - litellm.cache.async_get_cache(cache_key=preset_cache_key) - ) - 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 - elif isinstance(litellm.cache.cache, RedisSemanticCache) or isinstance( - litellm.cache.cache, RedisCache - ): - preset_cache_key = litellm.cache.get_cache_key(*args, **kwargs) - kwargs["preset_cache_key"] = ( - preset_cache_key # for streaming calls, we need to pass the preset_cache_key - ) - cached_result = await litellm.cache.async_get_cache(*args, **kwargs) - elif isinstance(litellm.cache.cache, QdrantSemanticCache): - preset_cache_key = litellm.cache.get_cache_key(*args, **kwargs) - kwargs["preset_cache_key"] = ( - preset_cache_key # for streaming calls, we need to pass the preset_cache_key - ) - cached_result = await litellm.cache.async_get_cache(*args, **kwargs) - else: # for s3 caching. [NOT RECOMMENDED IN PROD - this will slow down responses since boto3 is sync] - preset_cache_key = litellm.cache.get_cache_key(*args, **kwargs) - kwargs["preset_cache_key"] = ( - preset_cache_key # for streaming calls, we need to pass the preset_cache_key - ) - cached_result = litellm.cache.get_cache(*args, **kwargs) + + cached_result = await self._retrieve_from_cache( + call_type=call_type, + kwargs=kwargs, + args=args, + ) + if cached_result is not None and not isinstance(cached_result, list): print_verbose("Cache Hit!") cache_hit = True @@ -193,69 +168,16 @@ async def _async_get_cache( stream=kwargs.get("stream", False), ) call_type = original_function.__name__ - if call_type == CallTypes.acompletion.value and isinstance( - cached_result, dict - ): - if kwargs.get("stream", False) is True: - cached_result = convert_to_streaming_response_async( - response_object=cached_result, - ) - cached_result = CustomStreamWrapper( - completion_stream=cached_result, - model=model, - custom_llm_provider="cached_response", - logging_obj=logging_obj, - ) - else: - cached_result = convert_to_model_response_object( - response_object=cached_result, - model_response_object=ModelResponse(), - ) - if call_type == CallTypes.atext_completion.value and isinstance( - cached_result, dict - ): - if kwargs.get("stream", False) is True: - cached_result = convert_to_streaming_response_async( - response_object=cached_result, - ) - cached_result = CustomStreamWrapper( - completion_stream=cached_result, - model=model, - custom_llm_provider="cached_response", - logging_obj=logging_obj, - ) - else: - cached_result = TextCompletionResponse(**cached_result) - elif call_type == CallTypes.aembedding.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 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 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, - ) + + 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 kwargs.get("stream", False) is False: # LOG SUCCESS asyncio.create_task( @@ -378,6 +300,151 @@ async def _async_get_cache( final_embedding_cached_response=final_embedding_cached_response, ) + 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 + + cached_result: Optional[Any] = None + if call_type == CallTypes.aembedding.value and isinstance( + kwargs["input"], list + ): + tasks = [] + for idx, i in enumerate(kwargs["input"]): + preset_cache_key = litellm.cache.get_cache_key( + *args, **{**kwargs, "input": i} + ) + tasks.append(litellm.cache.async_get_cache(cache_key=preset_cache_key)) + 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: + preset_cache_key = litellm.cache.get_cache_key(*args, **kwargs) + kwargs["preset_cache_key"] = ( + preset_cache_key # for streaming calls, we need to pass the preset_cache_key + ) + if litellm.cache._supports_async() is True: + cached_result = await litellm.cache.async_get_cache(*args, **kwargs) + else: # for s3 caching. [NOT RECOMMENDED IN PROD - this will slow down responses since boto3 is sync] + cached_result = litellm.cache.get_cache(*args, **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[Any]: + """ + 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 ( + CustomStreamWrapper, + convert_to_model_response_object, + convert_to_streaming_response_async, + ) + + if call_type == CallTypes.acompletion.value and isinstance(cached_result, dict): + if kwargs.get("stream", False) is True: + cached_result = convert_to_streaming_response_async( + response_object=cached_result, + ) + cached_result = CustomStreamWrapper( + completion_stream=cached_result, + model=model, + custom_llm_provider="cached_response", + logging_obj=logging_obj, + ) + else: + cached_result = convert_to_model_response_object( + response_object=cached_result, + model_response_object=ModelResponse(), + ) + if call_type == CallTypes.atext_completion.value and isinstance( + cached_result, dict + ): + if kwargs.get("stream", False) is True: + cached_result = convert_to_streaming_response_async( + response_object=cached_result, + ) + cached_result = CustomStreamWrapper( + completion_stream=cached_result, + model=model, + custom_llm_provider="cached_response", + logging_obj=logging_obj, + ) + else: + cached_result = TextCompletionResponse(**cached_result) + elif call_type == CallTypes.aembedding.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 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 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, + ) + return cached_result + async def _async_set_cache( self, result: Any, @@ -438,3 +505,45 @@ async def _async_set_cache( asyncio.create_task( litellm.cache.async_add_cache(result, *args, **kwargs) ) + + 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] + ] = None + if ( + processed_chunk.choices[0].finish_reason is not None + ): # if it's the last chunk + self.async_streaming_chunks.append(processed_chunk) + try: + end_time: datetime.datetime = datetime.datetime.now() + complete_streaming_response = litellm.stream_chunk_builder( + self.async_streaming_chunks, + messages=self.request_kwargs.get("messages", None), + start_time=self.start_time, + end_time=end_time, + ) + except Exception as e: + verbose_logger.exception( + "Error occurred building stream chunk in success logging: {}".format( + str(e) + ) + ) + complete_streaming_response = None + else: + self.async_streaming_chunks.append(processed_chunk) + + # 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, + ) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 4bffa94b3c7f..1d4a3e3cc550 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -24,6 +24,7 @@ verbose_logger, ) from litellm.caching.caching import DualCache, InMemoryCache, S3Cache +from litellm.caching.caching_handler import LLMCachingHandler from litellm.cost_calculator import _select_model_name_for_cost_calc from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.integrations.custom_logger import CustomLogger @@ -271,6 +272,7 @@ def __init__( ## TIME TO FIRST TOKEN LOGGING ## self.completion_start_time: Optional[datetime.datetime] = None + self._llm_caching_handler: Optional[LLMCachingHandler] = None def process_dynamic_callbacks(self): """ @@ -1625,35 +1627,6 @@ async def async_success_handler( if kwargs.get("no-log", False) is True: print_verbose("no-log request, skipping logging") continue - if ( - callback == "cache" - and litellm.cache is not None - and self.model_call_details.get("litellm_params", {}).get( - "acompletion", False - ) - is True - ): - # set_cache once complete streaming response is built - print_verbose("async success_callback: reaches cache for logging!") - kwargs = self.model_call_details - if self.stream: - if "async_complete_streaming_response" not in kwargs: - print_verbose( - f"async success_callback: reaches cache for logging, there is no async_complete_streaming_response. Kwargs={kwargs}\n\n" - ) - pass - else: - print_verbose( - "async success_callback: reaches cache for logging, there is a async_complete_streaming_response. Adding to cache" - ) - result = kwargs["async_complete_streaming_response"] - # only add to cache once we have a complete streaming response - if litellm.cache is not None and not isinstance( - litellm.cache.cache, S3Cache - ): - await litellm.cache.async_add_cache(result, **kwargs) - else: - litellm.cache.add_cache(result, **kwargs) if callback == "openmeter" and openMeterLogger is not None: if self.stream is True: if ( diff --git a/litellm/llms/vertex_ai_and_google_ai_studio/context_caching/vertex_ai_context_caching.py b/litellm/llms/vertex_ai_and_google_ai_studio/context_caching/vertex_ai_context_caching.py index 2dafce6a98b6..e60a1705283e 100644 --- a/litellm/llms/vertex_ai_and_google_ai_studio/context_caching/vertex_ai_context_caching.py +++ b/litellm/llms/vertex_ai_and_google_ai_studio/context_caching/vertex_ai_context_caching.py @@ -4,7 +4,7 @@ import httpx import litellm -from litellm.caching.caching import Cache +from litellm.caching.caching import Cache, LiteLLMCacheType from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.OpenAI.openai import AllMessageValues @@ -22,7 +22,9 @@ transform_openai_messages_to_gemini_context_caching, ) -local_cache_obj = Cache(type="local") # only used for calling 'get_cache_key' function +local_cache_obj = Cache( + type=LiteLLMCacheType.LOCAL +) # only used for calling 'get_cache_key' function class ContextCachingEndpoints(VertexBase): diff --git a/litellm/types/caching.py b/litellm/types/caching.py new file mode 100644 index 000000000000..7fca4c041287 --- /dev/null +++ b/litellm/types/caching.py @@ -0,0 +1,25 @@ +from enum import Enum +from typing import Literal + + +class LiteLLMCacheType(str, Enum): + LOCAL = "local" + REDIS = "redis" + REDIS_SEMANTIC = "redis-semantic" + S3 = "s3" + DISK = "disk" + QDRANT_SEMANTIC = "qdrant-semantic" + + +CachingSupportedCallTypes = Literal[ + "completion", + "acompletion", + "embedding", + "aembedding", + "atranscription", + "transcription", + "atext_completion", + "text_completion", + "arerank", + "rerank", +] diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 2a36dd84d345..c3118b453b96 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -1418,17 +1418,3 @@ class StandardCallbackDynamicParams(TypedDict, total=False): # GCS dynamic params gcs_bucket_name: Optional[str] gcs_path_service_account: Optional[str] - - -CachingSupportedCallTypes = Literal[ - "completion", - "acompletion", - "embedding", - "aembedding", - "atranscription", - "transcription", - "atext_completion", - "text_completion", - "arerank", - "rerank", -] diff --git a/litellm/utils.py b/litellm/utils.py index a79a16a58f88..2457bdf4c3e9 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -58,8 +58,6 @@ import litellm.litellm_core_utils.json_validation_rule from litellm.caching.caching import DualCache from litellm.caching.caching_handler import CachingHandlerResponse, LLMCachingHandler - -_llm_caching_handler = LLMCachingHandler() from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.core_helpers import map_finish_reason from litellm.litellm_core_utils.exception_mapping_utils import ( @@ -73,6 +71,7 @@ ) from litellm.litellm_core_utils.llm_request_utils import _ensure_extra_body_is_safe from litellm.litellm_core_utils.redact_messages import ( + LiteLLMLoggingObject, redact_message_input_output_from_logging, ) from litellm.litellm_core_utils.token_counter import get_modified_max_tokens @@ -1095,7 +1094,14 @@ async def wrapper_async(*args, **kwargs): print_args_passed_to_litellm(original_function, args, kwargs) start_time = datetime.datetime.now() result = None - logging_obj = kwargs.get("litellm_logging_obj", None) + logging_obj: Optional[LiteLLMLoggingObject] = kwargs.get( + "litellm_logging_obj", None + ) + _llm_caching_handler: LLMCachingHandler = LLMCachingHandler( + original_function=original_function, + request_kwargs=kwargs, + start_time=start_time, + ) # only set litellm_call_id if its not in kwargs call_type = original_function.__name__ if "litellm_call_id" not in kwargs: @@ -1117,7 +1123,7 @@ async def wrapper_async(*args, **kwargs): original_function.__name__, rules_obj, start_time, *args, **kwargs ) kwargs["litellm_logging_obj"] = logging_obj - + logging_obj._llm_caching_handler = _llm_caching_handler # [OPTIONAL] CHECK BUDGET if litellm.max_budget: if litellm._current_cost > litellm.max_budget: @@ -6254,7 +6260,7 @@ def __init__( self.model = model self.make_call = make_call self.custom_llm_provider = custom_llm_provider - self.logging_obj = logging_obj + self.logging_obj: LiteLLMLoggingObject = logging_obj self.completion_stream = completion_stream self.sent_first_chunk = False self.sent_last_chunk = False @@ -8067,6 +8073,14 @@ async def __anext__(self): processed_chunk, cache_hit=cache_hit ) ) + + if self.logging_obj._llm_caching_handler is not None: + asyncio.create_task( + self.logging_obj._llm_caching_handler._add_streaming_response_to_cache( + processed_chunk=processed_chunk, + ) + ) + choice = processed_chunk.choices[0] if isinstance(choice, StreamingChoices): self.response_uptil_now += choice.delta.get("content", "") or ""