fix use caching lib

This commit is contained in:
Ishaan Jaff 2024-06-24 17:03:23 -07:00
parent 8a66e074ce
commit 974d92ff45

View file

@ -17,6 +17,7 @@ import traceback
from datetime import timedelta from datetime import timedelta
from typing import Any, BinaryIO, List, Literal, Optional, Union from typing import Any, BinaryIO, List, Literal, Optional, Union
from cachetools import Cache as CachetoolsCache
from openai._models import BaseModel as OpenAIObject from openai._models import BaseModel as OpenAIObject
import litellm import litellm
@ -70,7 +71,9 @@ class InMemoryCache(BaseCache):
this is done to prevent overuse of System RAM this is done to prevent overuse of System RAM
""" """
# if users don't provider one, use the default litellm cache # if users don't provider one, use the default litellm cache
self.cache_dict: dict = {} self.cache_dict: CachetoolsCache = CachetoolsCache(
maxsize=1000,
)
self.ttl_dict: dict = {} self.ttl_dict: dict = {}
self.default_ttl = default_ttl or 120.0 self.default_ttl = default_ttl or 120.0
self.last_cleaned = 0 # since this is in memory we need to periodically clean it up to not overuse the machines RAM self.last_cleaned = 0 # since this is in memory we need to periodically clean it up to not overuse the machines RAM
@ -83,8 +86,6 @@ class InMemoryCache(BaseCache):
async def async_set_cache(self, key, value, **kwargs): async def async_set_cache(self, key, value, **kwargs):
self.set_cache(key=key, value=value, **kwargs) self.set_cache(key=key, value=value, **kwargs)
if time.time() - self.last_cleaned > self.default_ttl:
asyncio.create_task(self.clean_up_in_memory_cache())
async def async_set_cache_pipeline(self, cache_list, ttl=None): async def async_set_cache_pipeline(self, cache_list, ttl=None):
for cache_key, cache_value in cache_list: for cache_key, cache_value in cache_list:
@ -93,10 +94,6 @@ class InMemoryCache(BaseCache):
else: else:
self.set_cache(key=cache_key, value=cache_value) self.set_cache(key=cache_key, value=cache_value)
if time.time() - self.last_cleaned > self.default_ttl:
asyncio.create_task(self.clean_up_in_memory_cache())
async def async_set_cache_sadd(self, key, value: List, ttl: Optional[float]): async def async_set_cache_sadd(self, key, value: List, ttl: Optional[float]):
""" """
Add value to set Add value to set
@ -108,7 +105,6 @@ class InMemoryCache(BaseCache):
self.set_cache(key, init_value, ttl=ttl) self.set_cache(key, init_value, ttl=ttl)
return value return value
def get_cache(self, key, **kwargs): def get_cache(self, key, **kwargs):
if key in self.cache_dict: if key in self.cache_dict:
if key in self.ttl_dict: if key in self.ttl_dict: