Merge pull request #722 from karvetskiy/fix-router-caching

Fix caching for Router

litellm/caching.py

@@ -11,6 +11,7 @@ import litellm
 import time
 import json
 
+
 def get_prompt(*args, **kwargs):
     # make this safe checks, it should not throw any exceptions
     if len(args) > 1:
@@ -23,20 +24,30 @@ def get_prompt(*args, **kwargs):
         return prompt
     return None
 
-class RedisCache():
+
+class BaseCache:
+    def set_cache(self, key, value, **kwargs):
+        raise NotImplementedError
+
+    def get_cache(self, key, **kwargs):
+        raise NotImplementedError
+
+
+class RedisCache(BaseCache):
     def __init__(self, host, port, password):
         import redis
         # if users don't provider one, use the default litellm cache
         self.redis_client = redis.Redis(host=host, port=port, password=password)
 
-    def set_cache(self, key, value):
+    def set_cache(self, key, value, **kwargs):
+        ttl = kwargs.get("ttl", None)
         try:
-            self.redis_client.set(key, str(value))
+            self.redis_client.set(name=key, value=str(value), ex=ttl)
         except Exception as e:
             # NON blocking - notify users Redis is throwing an exception
             print("LiteLLM Caching: Got exception from REDIS: ", e)
 
-    def get_cache(self, key):
+    def get_cache(self, key, **kwargs):
         try:
             # TODO convert this to a ModelResponse object
             cached_response = self.redis_client.get(key)
@@ -50,14 +61,17 @@ class RedisCache():
             # NON blocking - notify users Redis is throwing an exception
             print("LiteLLM Caching: Got exception from REDIS: ", e)
 
-class HostedCache():
+
-    def set_cache(self, key, value):
+class HostedCache(BaseCache):
+    def set_cache(self, key, value, **kwargs):
+        if "ttl" in kwargs:
+            print("LiteLLM Caching: TTL is not supported for hosted cache!")
         # make a post request to api.litellm.ai/set_cache
         import requests
         url = f"https://api.litellm.ai/set_cache?key={key}&value={str(value)}"
         requests.request("POST", url)  # post request to set this in the hosted litellm cache
 
-    def get_cache(self, key):
+    def get_cache(self, key, **kwargs):
         import requests
         url = f"https://api.litellm.ai/get_cache?key={key}"
         cached_response = requests.request("GET", url)
@@ -72,25 +86,29 @@ class HostedCache():
             except:
                 return cached_response
 
-class InMemoryCache():
+
+class InMemoryCache(BaseCache):
     def __init__(self):
         # if users don't provider one, use the default litellm cache
         self.cache_dict = {}
+        self.ttl_dict = {}
 
-    def set_cache(self, key, value):
+    def set_cache(self, key, value, **kwargs):
-        #print("in set cache for inmem")
         self.cache_dict[key] = value
-        #print(self.cache_dict)
+        if "ttl" in kwargs:
+            self.ttl_dict[key] = time.time() + kwargs["ttl"]
 
-    def get_cache(self, key):
+    def get_cache(self, key, **kwargs):
-        #print("in get cache for inmem")
         if key in self.cache_dict:
-            #print("got a cache hit")
+            if key in self.ttl_dict:
+                if time.time() > self.ttl_dict[key]:
+                    self.cache_dict.pop(key, None)
+                    return None
             return self.cache_dict[key]
-        #print("got a cache miss")
         return None
 
-class Cache():
+
+class Cache:
     def __init__(
             self,
             type="local",
@@ -201,12 +219,6 @@ class Cache():
             # print(cache_key)
             if cache_key is not None:
                 # print("adding to cache", cache_key, result)
-                self.cache.set_cache(cache_key, result)
+                self.cache.set_cache(cache_key, result, **kwargs)
         except:
             pass
-
-
-
-
-
-

litellm/router.py

@@ -24,6 +24,8 @@ class Router:
     """
     model_names: List = []
     cache_responses: bool = False
+    default_cache_time_seconds: int = 1 * 60 * 60  # 1 hour
+
     def __init__(self,
                  model_list: Optional[list] = None,
                  redis_host: Optional[str] = None,
@@ -133,7 +135,10 @@ class Router:
         Function LiteLLM submits a callback to after a successful
         completion. Purpose of this is ti update TPM/RPM usage per model
         """
-        model_name = kwargs.get('model', None)  # i.e. azure/gpt35turbo
+        model_name = kwargs.get('model', None)  # i.e. gpt35turbo
+        custom_llm_provider = kwargs.get("litellm_params", {}).get('custom_llm_provider', None)  # i.e. azure
+        if custom_llm_provider:
+            model_name = f"{custom_llm_provider}/{model_name}"
         total_tokens = completion_response['usage']['total_tokens']
         self._set_deployment_usage(model_name, total_tokens)
 
@@ -150,17 +155,9 @@ class Router:
             if item["model_name"] == model:
                 potential_deployments.append(item)
 
-        # set first model as current model
+        # set first model as current model to calculate token count
         deployment = potential_deployments[0]
 
-
-        # get model tpm, rpm limits
-        tpm = deployment["tpm"]
-        rpm = deployment["rpm"]
-
-        # get deployment current usage
-        current_tpm, current_rpm = self._get_deployment_usage(deployment_name=deployment["litellm_params"]["model"])
-
         # get encoding
         token_count = 0
         if messages is not None:
@@ -172,12 +169,10 @@ class Router:
                 input_text = input
             token_count = litellm.token_counter(model=deployment["model_name"], text=input_text)
 
-        # if at model limit, return lowest used
-        if current_tpm + token_count > tpm or current_rpm + 1 >= rpm:
         # -----------------------
         # Find lowest used model
         # ----------------------
-            lowest_tpm = float('inf')
+        lowest_tpm = float("inf")
         deployment = None
 
         # Go through all the models to get tpm, rpm
@@ -194,7 +189,7 @@ class Router:
 
         # if none, raise exception
         if deployment is None:
-                raise ValueError(f"No models available.")
+            raise ValueError("No models available.")
 
         # return model
         return deployment
@@ -213,26 +208,21 @@ class Router:
         # ------------
         # Return usage
         # ------------
-        tpm = self.cache.get_cache(tpm_key)
+        tpm = self.cache.get_cache(cache_key=tpm_key) or 0
-        rpm = self.cache.get_cache(rpm_key)
+        rpm = self.cache.get_cache(cache_key=rpm_key) or 0
-
-        if tpm is None:
-            tpm = 0
-        if rpm is None:
-            rpm = 0
 
         return int(tpm), int(rpm)
 
     def increment(self, key: str, increment_value: int):
         # get value
-        cached_value = self.cache.get_cache(key)
+        cached_value = self.cache.get_cache(cache_key=key)
         # update value
         try:
             cached_value = cached_value + increment_value
         except:
             cached_value = increment_value
         # save updated value
-        self.cache.add_cache(result=cached_value, cache_key=key)
+        self.cache.add_cache(result=cached_value, cache_key=key, ttl=self.default_cache_time_seconds)
 
     def _set_deployment_usage(
         self,

litellm/tests/test_caching.py

@@ -1,4 +1,5 @@
 import sys, os
+import time
 import traceback
 from dotenv import load_dotenv
 
@@ -344,7 +345,7 @@ def test_custom_redis_cache_with_key():
         pytest.fail(f"Error occurred:")
     litellm.cache = None
 
-test_custom_redis_cache_with_key()
+# test_custom_redis_cache_with_key()
 
 def test_hosted_cache():
     litellm.cache = Cache(type="hosted") # use api.litellm.ai for caching
@@ -364,3 +365,20 @@ def test_hosted_cache():
 
 # test_hosted_cache()
 
+
+def test_redis_cache_with_ttl():
+    cache = Cache(type="redis", host=os.environ['REDIS_HOST'], port=os.environ['REDIS_PORT'], password=os.environ['REDIS_PASSWORD'])
+    cache.add_cache(cache_key="test_key", result="test_value", ttl=1)
+    cached_value = cache.get_cache(cache_key="test_key")
+    assert cached_value == "test_value"
+    time.sleep(2)
+    assert cache.get_cache(cache_key="test_key") is None
+
+
+def test_in_memory_cache_with_ttl():
+    cache = Cache(type="local")
+    cache.add_cache(cache_key="test_key", result="test_value", ttl=1)
+    cached_value = cache.get_cache(cache_key="test_key")
+    assert cached_value == "test_value"
+    time.sleep(2)
+    assert cache.get_cache(cache_key="test_key") is None