phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-27 11:43:54 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

build(caching.py): move diskcache import inside class and add cache_dir argument to Cache

Browse source
This commit is contained in:
Antonio Loison 2024-05-10 12:04:54 +02:00
parent 79d774e330
commit dabf4cf8a8
1 changed files with 8 additions and 3 deletions
Download patch file Download diff file Expand all files Collapse all files

11
litellm/caching.py
View file

@ -1491,6 +1491,7 @@ class Cache:
redis_semantic_cache_use_async=False, redis_semantic_cache_use_async=False,
redis_semantic_cache_embedding_model="text-embedding-ada-002", redis_semantic_cache_embedding_model="text-embedding-ada-002",
redis_flush_size=None, redis_flush_size=None,
disk_cache_dir=None,
**kwargs, **kwargs,
): ):
""" """
@ -1544,7 +1545,7 @@ class Cache:
**kwargs, **kwargs,
) )
elif type == "disk": elif type == "disk":
self.cache = DiskCache() self.cache = DiskCache(disk_cache_dir=disk_cache_dir)
if "cache" not in litellm.input_callback: if "cache" not in litellm.input_callback:
litellm.input_callback.append("cache") litellm.input_callback.append("cache")
if "cache" not in litellm.success_callback: if "cache" not in litellm.success_callback:
@ -1917,10 +1918,14 @@ class Cache:
class DiskCache(BaseCache): class DiskCache(BaseCache):
def __init__(self, cache_dir: str = ".litellm_cache"): def __init__(self, disk_cache_dir: Optional[str] = None):
import diskcache as dc import diskcache as dc
# if users don't provider one, use the default litellm cache # if users don't provider one, use the default litellm cache
self.disk_cache = dc.Cache(cache_dir) if disk_cache_dir is None:
self.disk_cache = dc.Cache(".litellm_cache")
else:
self.disk_cache = dc.Cache(disk_cache_dir)
def set_cache(self, key, value, **kwargs): def set_cache(self, key, value, **kwargs):
print_verbose("DiskCache: set_cache") print_verbose("DiskCache: set_cache")