fix(proxy_server.py): allow user to set in-memory + redis ttl

addresses - https://github.com/BerriAI/litellm/issues/2700
This commit is contained in:
Krrish Dholakia 2024-04-01 19:13:23 -07:00
parent ceabf726b0
commit 203e2776f8
3 changed files with 22 additions and 2 deletions

View file

@ -76,6 +76,8 @@ caching_with_models: bool = (
cache: Optional[Cache] = ( cache: Optional[Cache] = (
None # cache object <- use this - https://docs.litellm.ai/docs/caching None # cache object <- use this - https://docs.litellm.ai/docs/caching
) )
default_in_memory_ttl: Optional[float] = None
default_redis_ttl: Optional[float] = None
model_alias_map: Dict[str, str] = {} model_alias_map: Dict[str, str] = {}
model_group_alias_map: Dict[str, str] = {} model_group_alias_map: Dict[str, str] = {}
max_budget: float = 0.0 # set the max budget across all providers max_budget: float = 0.0 # set the max budget across all providers

View file

@ -828,8 +828,10 @@ class DualCache(BaseCache):
# If redis_cache is not provided, use the default RedisCache # If redis_cache is not provided, use the default RedisCache
self.redis_cache = redis_cache self.redis_cache = redis_cache
self.default_in_memory_ttl = default_in_memory_ttl self.default_in_memory_ttl = (
self.default_redis_ttl = default_redis_ttl default_in_memory_ttl or litellm.default_in_memory_ttl
)
self.default_redis_ttl = default_redis_ttl or litellm.default_redis_ttl
def set_cache(self, key, value, local_only: bool = False, **kwargs): def set_cache(self, key, value, local_only: bool = False, **kwargs):
# Update both Redis and in-memory cache # Update both Redis and in-memory cache
@ -939,6 +941,8 @@ class Cache:
password: Optional[str] = None, password: Optional[str] = None,
namespace: Optional[str] = None, namespace: Optional[str] = None,
ttl: Optional[float] = None, ttl: Optional[float] = None,
default_in_memory_ttl: Optional[float] = None,
default_in_redis_ttl: Optional[float] = None,
similarity_threshold: Optional[float] = None, similarity_threshold: Optional[float] = None,
supported_call_types: Optional[ supported_call_types: Optional[
List[ List[
@ -1038,6 +1042,14 @@ class Cache:
self.redis_flush_size = redis_flush_size self.redis_flush_size = redis_flush_size
self.ttl = ttl self.ttl = ttl
if self.type == "local" and default_in_memory_ttl is not None:
self.ttl = default_in_memory_ttl
if (
self.type == "redis" or self.type == "redis-semantic"
) and default_in_redis_ttl is not None:
self.ttl = default_in_redis_ttl
if self.namespace is not None and isinstance(self.cache, RedisCache): if self.namespace is not None and isinstance(self.cache, RedisCache):
self.cache.namespace = self.namespace self.cache.namespace = self.namespace

View file

@ -1914,6 +1914,12 @@ class ProxyConfig:
global redis_usage_cache global redis_usage_cache
from litellm import Cache from litellm import Cache
if "default_in_memory_ttl" in cache_params:
litellm.default_in_memory_ttl = cache_params["default_in_memory_ttl"]
if "default_in_redis_ttl" in cache_params:
litellm.default_redis_ttl = cache_params["default_in_redis_ttl"]
litellm.cache = Cache(**cache_params) litellm.cache = Cache(**cache_params)
if litellm.cache is not None and isinstance(litellm.cache.cache, RedisCache): if litellm.cache is not None and isinstance(litellm.cache.cache, RedisCache):