diff --git a/litellm/caching/caching.py b/litellm/caching/caching.py index 6a7c93e3fe..e407e0ce00 100644 --- a/litellm/caching/caching.py +++ b/litellm/caching/caching.py @@ -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: diff --git a/litellm/caching/disk_cache.py b/litellm/caching/disk_cache.py index 413ac2932d..65fdedd3ea 100644 --- a/litellm/caching/disk_cache.py +++ b/litellm/caching/disk_cache.py @@ -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: