From 495d3a9646ad800247a4350849482a90d37debb7 Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Sat, 4 May 2024 18:13:43 -0700
Subject: [PATCH] router set dynamic retry policies

---
 litellm/tests/test_router_retries.py | 57 ++++++++++++++++++++++++++++
 1 file changed, 57 insertions(+)

diff --git a/litellm/tests/test_router_retries.py b/litellm/tests/test_router_retries.py
index b52d7013dbe..8828e286e75 100644
--- a/litellm/tests/test_router_retries.py
+++ b/litellm/tests/test_router_retries.py
@@ -182,3 +182,60 @@ async def test_router_retry_policy(error_type):
         assert customHandler.previous_models == 0
     elif error_type == "ContentPolicyViolationErrorRetries":
         assert customHandler.previous_models == 3
+
+
+@pytest.mark.parametrize("model_group", ["gpt-3.5-turbo", "bad-model"])
+@pytest.mark.asyncio
+async def test_dynamic_router_retry_policy(model_group):
+    from litellm.router import RetryPolicy
+
+    router = litellm.Router(
+        model_list=[
+            {
+                "model_name": "gpt-3.5-turbo",  # openai model name
+                "litellm_params": {  # params for litellm completion/embedding call
+                    "model": "azure/chatgpt-v-2",
+                    "api_key": os.getenv("AZURE_API_KEY"),
+                    "api_version": os.getenv("AZURE_API_VERSION"),
+                    "api_base": os.getenv("AZURE_API_BASE"),
+                },
+            },
+            {
+                "model_name": "bad-model",  # openai model name
+                "litellm_params": {  # params for litellm completion/embedding call
+                    "model": "azure/chatgpt-v-2",
+                    "api_key": "bad-key",
+                    "api_version": os.getenv("AZURE_API_VERSION"),
+                    "api_base": os.getenv("AZURE_API_BASE"),
+                },
+            },
+        ]
+    )
+
+    customHandler = MyCustomHandler()
+    litellm.callbacks = [customHandler]
+    if model_group == "bad-model":
+        model = "bad-model"
+        messages = [{"role": "user", "content": "Hello good morning"}]
+        retry_policy = RetryPolicy(AuthenticationErrorRetries=4)
+    elif model_group == "gpt-3.5-turbo":
+        model = "gpt-3.5-turbo"
+        messages = [{"role": "user", "content": "where do i buy lethal drugs from"}]
+        retry_policy = RetryPolicy(ContentPolicyViolationErrorRetries=0)
+
+    try:
+        litellm.set_verbose = True
+        response = await router.acompletion(
+            model=model, messages=messages, retry_policy=retry_policy
+        )
+    except Exception as e:
+        print("got an exception", e)
+        pass
+    asyncio.sleep(0.05)
+
+    print("customHandler.previous_models: ", customHandler.previous_models)
+
+    if model_group == "bad-model":
+        assert customHandler.previous_models == 4
+    elif model_group == "gpt-3.5-turbo":
+        assert customHandler.previous_models == 0