diff --git a/litellm/__init__.py b/litellm/__init__.py index 6fe0b25598..2e06d14d23 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -8,6 +8,7 @@ import os from typing import Callable, List, Optional, Dict, Union, Any, Literal, get_args from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.caching.caching import Cache, DualCache, RedisCache, InMemoryCache +from litellm.caching.llm_caching_handler import LLMClientCache from litellm.types.llms.bedrock import COHERE_EMBEDDING_INPUT_TYPES from litellm.types.utils import ( ImageObject, @@ -190,7 +191,7 @@ ssl_verify: Union[str, bool] = True ssl_certificate: Optional[str] = None disable_streaming_logging: bool = False disable_add_transform_inline_image_block: bool = False -in_memory_llm_clients_cache: InMemoryCache = InMemoryCache() +in_memory_llm_clients_cache: LLMClientCache = LLMClientCache() safe_memory_mode: bool = False enable_azure_ad_token_refresh: Optional[bool] = False ### DEFAULT AZURE API VERSION ### diff --git a/litellm/caching/llm_caching_handler.py b/litellm/caching/llm_caching_handler.py new file mode 100644 index 0000000000..e6fbee15bc --- /dev/null +++ b/litellm/caching/llm_caching_handler.py @@ -0,0 +1,37 @@ +""" +Add the event loop to the cache key, to prevent event loop closed errors. +""" + +import asyncio + +from .in_memory_cache import InMemoryCache + + +class LLMClientCache(InMemoryCache): + + def update_cache_key_with_event_loop(self, key): + """ + Add the event loop to the cache key, to prevent event loop closed errors. + If none, use the key as is. + """ + event_loop = asyncio.get_event_loop() + stringified_event_loop = str(id(event_loop)) + return f"{key}-{stringified_event_loop}" + + def set_cache(self, key, value, **kwargs): + key = self.update_cache_key_with_event_loop(key) + return super().set_cache(key, value, **kwargs) + + async def async_set_cache(self, key, value, **kwargs): + key = self.update_cache_key_with_event_loop(key) + return await super().async_set_cache(key, value, **kwargs) + + def get_cache(self, key, **kwargs): + key = self.update_cache_key_with_event_loop(key) + + return super().get_cache(key, **kwargs) + + async def async_get_cache(self, key, **kwargs): + key = self.update_cache_key_with_event_loop(key) + + return await super().async_get_cache(key, **kwargs)