From 591a0a376e117f371c56d0eb8e8ca313e3617ef2 Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Mon, 25 Mar 2024 13:40:17 -0700
Subject: [PATCH] fix(caching.py): support default ttl for caching

---
 litellm/caching.py            |  6 ++++++
 litellm/tests/test_caching.py | 17 +++++++++++++++++
 2 files changed, 23 insertions(+)

diff --git a/litellm/caching.py b/litellm/caching.py
index 5a9008342..fcebeb91d 100644
--- a/litellm/caching.py
+++ b/litellm/caching.py
@@ -874,6 +874,7 @@ class Cache:
         port: Optional[str] = None,
         password: Optional[str] = None,
         namespace: Optional[str] = None,
+        ttl: Optional[float] = None,
         similarity_threshold: Optional[float] = None,
         supported_call_types: Optional[
             List[
@@ -967,6 +968,7 @@ class Cache:
         self.supported_call_types = supported_call_types  # default to ["completion", "acompletion", "embedding", "aembedding"]
         self.type = type
         self.namespace = namespace
+        self.ttl = ttl
 
     def get_cache_key(self, *args, **kwargs):
         """
@@ -1206,6 +1208,9 @@ class Cache:
                 if isinstance(result, OpenAIObject):
                     result = result.model_dump_json()
 
+                ## DEFAULT TTL ##
+                if self.ttl is not None:
+                    kwargs["ttl"] = self.ttl
                 ## Get Cache-Controls ##
                 if kwargs.get("cache", None) is not None and isinstance(
                     kwargs.get("cache"), dict
@@ -1213,6 +1218,7 @@ class Cache:
                     for k, v in kwargs.get("cache").items():
                         if k == "ttl":
                             kwargs["ttl"] = v
+
                 cached_data = {"timestamp": time.time(), "response": result}
                 return cache_key, cached_data, kwargs
             else:
diff --git a/litellm/tests/test_caching.py b/litellm/tests/test_caching.py
index 678e8ab59..4bc5d1d5a 100644
--- a/litellm/tests/test_caching.py
+++ b/litellm/tests/test_caching.py
@@ -116,6 +116,23 @@ def test_caching_with_ttl():
         pytest.fail(f"Error occurred: {e}")
 
 
+def test_caching_with_default_ttl():
+    try:
+        litellm.set_verbose = True
+        litellm.cache = Cache(ttl=0)
+        response1 = completion(model="gpt-3.5-turbo", messages=messages, caching=True)
+        response2 = completion(model="gpt-3.5-turbo", messages=messages, caching=True)
+        print(f"response1: {response1}")
+        print(f"response2: {response2}")
+        litellm.cache = None  # disable cache
+        litellm.success_callback = []
+        litellm._async_success_callback = []
+        assert response2["id"] != response1["id"]
+    except Exception as e:
+        print(f"error occurred: {traceback.format_exc()}")
+        pytest.fail(f"Error occurred: {e}")
+
+
 def test_caching_with_cache_controls():
     try:
         litellm.set_verbose = True