From f0b6b9dce2a88cb2189758ae5a5120ae8b2ce499 Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Tue, 26 Dec 2023 17:22:31 +0530
Subject: [PATCH] fix(main.py): support ttl being set for completion,
 embedding, image generation calls

---
 litellm/main.py                      |  3 ++
 litellm/tests/test_caching.py        | 24 +++++++++-
 litellm/tests/test_router_caching.py | 65 ++++++++++++++++++++++++++++
 3 files changed, 91 insertions(+), 1 deletion(-)

diff --git a/litellm/main.py b/litellm/main.py
index 65ece212af9..00d86ca20d5 100644
--- a/litellm/main.py
+++ b/litellm/main.py
@@ -466,6 +466,7 @@ def completion(
         "proxy_server_request",
         "preset_cache_key",
         "caching_groups",
+        "ttl",
     ]
     default_params = openai_params + litellm_params
     non_default_params = {
@@ -2158,6 +2159,7 @@ def embedding(
         "model_info",
         "preset_cache_key",
         "caching_groups",
+        "ttl",
     ]
     default_params = openai_params + litellm_params
     non_default_params = {
@@ -2787,6 +2789,7 @@ def image_generation(
         "model_info",
         "preset_cache_key",
         "caching_groups",
+        "ttl",
     ]
     default_params = openai_params + litellm_params
     non_default_params = {
diff --git a/litellm/tests/test_caching.py b/litellm/tests/test_caching.py
index 081f71ebb99..6736814b82a 100644
--- a/litellm/tests/test_caching.py
+++ b/litellm/tests/test_caching.py
@@ -58,6 +58,28 @@ def test_caching_v2():  # test in memory cache
 # test_caching_v2()
 
 
+def test_caching_with_ttl():
+    try:
+        litellm.set_verbose = True
+        litellm.cache = Cache()
+        response1 = completion(
+            model="gpt-3.5-turbo", messages=messages, caching=True, ttl=0
+        )
+        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["choices"][0]["message"]["content"]
+            != response1["choices"][0]["message"]["content"]
+        )
+    except Exception as e:
+        print(f"error occurred: {traceback.format_exc()}")
+        pytest.fail(f"Error occurred: {e}")
+
+
 def test_caching_with_models_v2():
     messages = [
         {"role": "user", "content": "who is ishaan CTO of litellm from litellm 2023"}
@@ -724,7 +746,7 @@ def test_get_cache_key():
         pytest.fail(f"Error occurred:", e)
 
 
-test_get_cache_key()
+# test_get_cache_key()
 
 # test_custom_redis_cache_params()
 
diff --git a/litellm/tests/test_router_caching.py b/litellm/tests/test_router_caching.py
index 67c263aa2eb..155c703ca76 100644
--- a/litellm/tests/test_router_caching.py
+++ b/litellm/tests/test_router_caching.py
@@ -80,6 +80,71 @@ async def test_acompletion_caching_on_router():
         pytest.fail(f"Error occurred: {e}")
 
 
+@pytest.mark.asyncio
+async def test_acompletion_caching_with_ttl_on_router():
+    # tests acompletion + caching on router
+    try:
+        litellm.set_verbose = True
+        model_list = [
+            {
+                "model_name": "gpt-3.5-turbo",
+                "litellm_params": {
+                    "model": "gpt-3.5-turbo-0613",
+                    "api_key": os.getenv("OPENAI_API_KEY"),
+                },
+                "tpm": 100000,
+                "rpm": 10000,
+            },
+            {
+                "model_name": "gpt-3.5-turbo",
+                "litellm_params": {
+                    "model": "azure/chatgpt-v-2",
+                    "api_key": os.getenv("AZURE_API_KEY"),
+                    "api_base": os.getenv("AZURE_API_BASE"),
+                    "api_version": os.getenv("AZURE_API_VERSION"),
+                },
+                "tpm": 100000,
+                "rpm": 10000,
+            },
+        ]
+
+        messages = [
+            {"role": "user", "content": f"write a one sentence poem {time.time()}?"}
+        ]
+        start_time = time.time()
+        router = Router(
+            model_list=model_list,
+            redis_host=os.environ["REDIS_HOST"],
+            redis_password=os.environ["REDIS_PASSWORD"],
+            redis_port=os.environ["REDIS_PORT"],
+            cache_responses=True,
+            timeout=30,
+            routing_strategy="simple-shuffle",
+        )
+        response1 = await router.acompletion(
+            model="gpt-3.5-turbo", messages=messages, temperature=1, ttl=0
+        )
+        print(f"response1: {response1}")
+        await asyncio.sleep(1)  # add cache is async, async sleep for cache to get set
+        response2 = await router.acompletion(
+            model="gpt-3.5-turbo", messages=messages, temperature=1, ttl=0
+        )
+        print(f"response2: {response2}")
+        assert response1.id != response2.id
+        assert len(response1.choices[0].message.content) > 0
+        assert (
+            response1.choices[0].message.content != response2.choices[0].message.content
+        )
+        router.reset()
+    except litellm.Timeout as e:
+        end_time = time.time()
+        print(f"timeout error occurred: {end_time - start_time}")
+        pass
+    except Exception as e:
+        traceback.print_exc()
+        pytest.fail(f"Error occurred: {e}")
+
+
 @pytest.mark.asyncio
 async def test_acompletion_caching_on_router_caching_groups():
     # tests acompletion + caching on router