From a9cbf2209c500ffdbff89a94fa6ef4a8644d0c6a Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Tue, 11 Jun 2024 17:51:29 -0700 Subject: [PATCH] fix(utils.py): support dynamic api key for azure_ai route --- litellm/tests/test_completion.py | 6 ++---- litellm/utils.py | 6 +++--- 2 files changed, 5 insertions(+), 7 deletions(-) diff --git a/litellm/tests/test_completion.py b/litellm/tests/test_completion.py index 3b6b507d6..8a7557c35 100644 --- a/litellm/tests/test_completion.py +++ b/litellm/tests/test_completion.py @@ -744,10 +744,7 @@ def test_completion_claude_3_function_plus_image(): @pytest.mark.parametrize( "provider", - [ - "azure", - "azure_ai" - ], + ["azure", "azure_ai"], ) def test_completion_azure_mistral_large_function_calling(provider): """ @@ -780,6 +777,7 @@ def test_completion_azure_mistral_large_function_calling(provider): "content": "What's the weather like in Boston today in Fahrenheit?", } ] + response = completion( model="{}/mistral-large-latest".format(provider), api_base=os.getenv("AZURE_MISTRAL_API_BASE"), diff --git a/litellm/utils.py b/litellm/utils.py index 9ad3b8e46..329f3185d 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -6581,6 +6581,9 @@ def get_llm_provider( or get_secret("FIREWORKSAI_API_KEY") or get_secret("FIREWORKS_AI_TOKEN") ) + elif custom_llm_provider == "azure_ai": + api_base = api_base or get_secret("AZURE_AI_API_BASE") # type: ignore + dynamic_api_key = api_key or get_secret("AZURE_AI_API_KEY") elif custom_llm_provider == "mistral": # mistral is openai compatible, we just need to set this to custom_openai and have the api_base be https://api.mistral.ai api_base = ( @@ -6599,9 +6602,6 @@ def get_llm_provider( or get_secret("MISTRAL_AZURE_API_KEY") # for Azure AI Mistral or get_secret("MISTRAL_API_KEY") ) - elif custom_llm_provider == "azure_ai": - api_base = api_base or get_secret("AZURE_AI_API_BASE") # type: ignore - dynamic_api_key = get_secret("AZURE_AI_API_KEY") elif custom_llm_provider == "voyage": # voyage is openai compatible, we just need to set this to custom_openai and have the api_base be https://api.voyageai.com/v1 api_base = "https://api.voyageai.com/v1"