test(test_router_fallbacks.py): add testing for sync streaming fallbacks

ventz · Dec 27, 2023 · 6d63c00 · 6d63c00
1 parent c9fdbaf
commit 6d63c00
Showing 1 changed file with 79 additions and 0 deletions.
diff --git a/litellm/tests/test_router_fallbacks.py b/litellm/tests/test_router_fallbacks.py
@@ -482,3 +482,82 @@ async def test_async_fallbacks_streaming():
         pytest.fail(f"An exception occurred: {e}")
     finally:
         router.reset()
+
+
+def test_sync_fallbacks_streaming():
+    try:
+        model_list = [
+            {  # list of model deployments
+                "model_name": "azure/gpt-3.5-turbo",  # 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"),
+                },
+                "tpm": 240000,
+                "rpm": 1800,
+            },
+            {  # list of model deployments
+                "model_name": "azure/gpt-3.5-turbo-context-fallback",  # 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"),
+                },
+                "tpm": 240000,
+                "rpm": 1800,
+            },
+            {
+                "model_name": "azure/gpt-3.5-turbo",  # openai model name
+                "litellm_params": {  # params for litellm completion/embedding call
+                    "model": "azure/chatgpt-functioncalling",
+                    "api_key": "bad-key",
+                    "api_version": os.getenv("AZURE_API_VERSION"),
+                    "api_base": os.getenv("AZURE_API_BASE"),
+                },
+                "tpm": 240000,
+                "rpm": 1800,
+            },
+            {
+                "model_name": "gpt-3.5-turbo",  # openai model name
+                "litellm_params": {  # params for litellm completion/embedding call
+                    "model": "gpt-3.5-turbo",
+                    "api_key": os.getenv("OPENAI_API_KEY"),
+                },
+                "tpm": 1000000,
+                "rpm": 9000,
+            },
+            {
+                "model_name": "gpt-3.5-turbo-16k",  # openai model name
+                "litellm_params": {  # params for litellm completion/embedding call
+                    "model": "gpt-3.5-turbo-16k",
+                    "api_key": os.getenv("OPENAI_API_KEY"),
+                },
+                "tpm": 1000000,
+                "rpm": 9000,
+            },
+        ]
+
+        litellm.set_verbose = True
+        customHandler = MyCustomHandler()
+        litellm.callbacks = [customHandler]
+        router = Router(
+            model_list=model_list,
+            fallbacks=[{"azure/gpt-3.5-turbo": ["gpt-3.5-turbo"]}],
+            context_window_fallbacks=[
+                {"azure/gpt-3.5-turbo-context-fallback": ["gpt-3.5-turbo-16k"]},
+                {"gpt-3.5-turbo": ["gpt-3.5-turbo-16k"]},
+            ],
+            set_verbose=False,
+        )
+        response = router.completion(**kwargs, stream=True)
+        print(f"response: {response}")
+        time.sleep(0.05)  # allow a delay as success_callbacks are on a separate thread
+        assert customHandler.previous_models == 1  # 0 retries, 1 fallback
+
+        print("Passed ! Test router_fallbacks: test_sync_fallbacks()")
+        router.reset()
+    except Exception as e:
+        print(e)