feat(router.py): add support for retry/fallbacks for async embedding calls

litellm/tests/test_router_fallbacks.py

@@ -227,6 +227,60 @@ async def test_async_fallbacks():
 # test_async_fallbacks()
 
 
+@pytest.mark.asyncio
+async def test_async_fallbacks_embeddings():
+    litellm.set_verbose = False
+    model_list = [
+        {  # list of model deployments
+            "model_name": "bad-azure-embedding-model",  # openai model name
+            "litellm_params": {  # params for litellm completion/embedding call
+                "model": "azure/azure-embedding-model",
+                "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": "good-azure-embedding-model",  # openai model name
+            "litellm_params": {  # params for litellm completion/embedding call
+                "model": "azure/azure-embedding-model",
+                "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,
+        },
+    ]
+
+    router = Router(
+        model_list=model_list,
+        fallbacks=[{"bad-azure-embedding-model": ["good-azure-embedding-model"]}],
+        set_verbose=False,
+    )
+    customHandler = MyCustomHandler()
+    litellm.callbacks = [customHandler]
+    user_message = "Hello, how are you?"
+    input = [user_message]
+    try:
+        kwargs = {"model": "bad-azure-embedding-model", "input": input}
+        response = await router.aembedding(**kwargs)
+        print(f"customHandler.previous_models: {customHandler.previous_models}")
+        await asyncio.sleep(
+            0.05
+        )  # allow a delay as success_callbacks are on a separate thread
+        assert customHandler.previous_models == 1  # 0 retries, 1 fallback
+        router.reset()
+    except litellm.Timeout as e:
+        pass
+    except Exception as e:
+        pytest.fail(f"An exception occurred: {e}")
+    finally:
+        router.reset()
+
+
 def test_dynamic_fallbacks_sync():
     """
     Allow setting the fallback in the router.completion() call.