This commit is contained in:
Joseph Chang 2025-04-24 00:57:22 -07:00 committed by GitHub
commit e811e8fbbc
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 22 additions and 4 deletions

View file

@ -94,6 +94,8 @@ class Cache:
redis_flush_size: Optional[int] = None,
redis_startup_nodes: Optional[List] = None,
disk_cache_dir: Optional[str] = None,
disk_cache_size_limit: Optional[float] = None,
disk_cache_cull_limit: Optional[int] = None,
qdrant_api_base: Optional[str] = None,
qdrant_api_key: Optional[str] = None,
qdrant_collection_name: Optional[str] = None,
@ -124,6 +126,8 @@ class Cache:
# Disk Cache Args
disk_cache_dir (str, optional): The directory for the disk cache. Defaults to None.
disk_cache_size_limit (int, optional): Default one gigabyte. The maximum on-disk size of the cache.
disk_cache_cull_limit (int, optional): Default ten. The maximum number of keys to cull when adding a new item. Set to zero to disable automatic culling. Some systems may disable automatic culling in exchange for a cron-like job that regularly calls cull in a separate process.
# S3 Cache Args
s3_bucket_name (str, optional): The bucket name for the s3 cache. Defaults to None.
@ -202,7 +206,11 @@ class Cache:
**kwargs,
)
elif type == LiteLLMCacheType.DISK:
self.cache = DiskCache(disk_cache_dir=disk_cache_dir)
self.cache = DiskCache(
disk_cache_dir=disk_cache_dir,
disk_cache_size_limit=disk_cache_size_limit,
disk_cache_cull_limit=disk_cache_cull_limit
)
if "cache" not in litellm.input_callback:
litellm.input_callback.append("cache")
if "cache" not in litellm.success_callback:

View file

@ -12,14 +12,24 @@ else:
class DiskCache(BaseCache):
def __init__(self, disk_cache_dir: Optional[str] = None):
def __init__(self,
disk_cache_dir: Optional[str] = None,
disk_cache_size_limit: Optional[float] = None,
disk_cache_cull_limit: Optional[int] = None,
) :
import diskcache as dc
# if users don't provider one, use the default litellm cache
disk_cache_params = {}
if disk_cache_size_limit is not None:
disk_cache_params["size_limit"] = disk_cache_size_limit
if disk_cache_cull_limit is not None:
disk_cache_params["cull_limit"] = disk_cache_cull_limit
if disk_cache_dir is None:
self.disk_cache = dc.Cache(".litellm_cache")
disk_cache_params["directory"] = ".litellm_cache"
else:
self.disk_cache = dc.Cache(disk_cache_dir)
disk_cache_params["directory"] = disk_cache_dir
self.disk_cache = dc.Cache(**disk_cache_params)
def set_cache(self, key, value, **kwargs):
if "ttl" in kwargs: