fix(utils.py): support dynamic api key for azure_ai route

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"),

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"