diff --git a/litellm/router.py b/litellm/router.py index 60d057c27..1798f75b9 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -327,7 +327,7 @@ class Router: data["model"] = original_model_string[:index_of_model_id] else: data["model"] = original_model_string - model_client = deployment.get("async_client", None) + model_client = deployment.get("client", None) return await litellm.aembedding(**{**data, "input": input, "caching": self.cache_responses, "client": model_client, **kwargs}) @@ -830,6 +830,7 @@ class Router: or custom_llm_provider == "openai" or custom_llm_provider == "azure" or "ft:gpt-3.5-turbo" in model_name + or model_name in litellm.open_ai_embedding_models ): # glorified / complicated reading of configs # user can pass vars directly or they can pas os.environ/AZURE_API_KEY, in which case we will read the env diff --git a/litellm/tests/test_router.py b/litellm/tests/test_router.py index e6afe3265..8024c9dd2 100644 --- a/litellm/tests/test_router.py +++ b/litellm/tests/test_router.py @@ -299,3 +299,34 @@ def test_aembedding_on_router(): traceback.print_exc() pytest.fail(f"Error occurred: {e}") # test_aembedding_on_router() + + +def test_azure_aembedding_on_router(): + litellm.set_verbose = True + try: + model_list = [ + { + "model_name": "text-embedding-ada-002", + "litellm_params": { + "model": "azure/azure-embedding-model", + "api_key":os.environ['AZURE_API_KEY'], + "api_base": os.environ['AZURE_API_BASE'] + }, + "tpm": 100000, + "rpm": 10000, + }, + ] + + async def embedding_call(): + router = Router(model_list=model_list) + response = await router.aembedding( + model="text-embedding-ada-002", + input=["good morning from litellm"] + ) + print(response) + router.reset() + asyncio.run(embedding_call()) + except Exception as e: + traceback.print_exc() + pytest.fail(f"Error occurred: {e}") +# test_azure_aembedding_on_router() \ No newline at end of file