diff --git a/litellm/caching.py b/litellm/caching.py index 78b4bd2708..b6921bac89 100644 --- a/litellm/caching.py +++ b/litellm/caching.py @@ -64,21 +64,53 @@ class BaseCache: class InMemoryCache(BaseCache): - def __init__(self, max_size_in_memory: Optional[int] = 200): + def __init__( + self, + max_size_in_memory: Optional[int] = 200, + default_ttl: Optional[ + int + ] = 300, # default ttl is 5 minutes. At maximum litellm rate limiting logic requires objects to be in memory for 1 minute + ): """ max_size_in_memory [int]: Maximum number of items in cache. done to prevent memory leaks. Use 200 items as a default """ - from cachetools import LRUCache + self.max_size_in_memory = ( + max_size_in_memory or 200 + ) # set an upper bound of 200 items in-memory + self.default_ttl = default_ttl or 300 - self.max_size_in_memory = max_size_in_memory or 200 - self.cache_dict: LRUCache = LRUCache(maxsize=self.max_size_in_memory) - self.ttl_dict: LRUCache = LRUCache(maxsize=self.max_size_in_memory) + # in-memory cache + self.cache_dict: dict = {} + self.ttl_dict: dict = {} + + def evict_cache(self): + """ + Eviction policy: + - check if any items in ttl_dict are expired -> remove them from ttl_dict and cache_dict + + + This guarantees the following: + - 1. When item ttl not set: At minimumm each item will remain in memory for 5 minutes + - 2. When ttl is set: the item will remain in memory for at least that amount of time + - 3. the size of in-memory cache is bounded + + """ + for key in list(self.ttl_dict.keys()): + if time.time() > self.ttl_dict[key]: + self.cache_dict.pop(key, None) + self.ttl_dict.pop(key, None) def set_cache(self, key, value, **kwargs): print_verbose("InMemoryCache: set_cache") + if len(self.cache_dict) >= self.max_size_in_memory: + # only evict when cache is full + self.evict_cache() + self.cache_dict[key] = value if "ttl" in kwargs: self.ttl_dict[key] = time.time() + kwargs["ttl"] + else: + self.ttl_dict[key] = time.time() + self.default_ttl async def async_set_cache(self, key, value, **kwargs): self.set_cache(key=key, value=value, **kwargs)