diff --git a/litellm/main.py b/litellm/main.py index 6320858167cf..e8022f464f55 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -692,9 +692,9 @@ def completion( or get_secret("AZURE_API_KEY") ) - azure_ad_token = optional_params.pop("azure_ad_token", None) or get_secret( - "AZURE_AD_TOKEN" - ) + azure_ad_token = optional_params.get("extra_body", {}).pop( + "azure_ad_token", None + ) or get_secret("AZURE_AD_TOKEN") headers = headers or litellm.headers diff --git a/litellm/tests/test_completion.py b/litellm/tests/test_completion.py index 71884bde736b..ab8673bca1fe 100644 --- a/litellm/tests/test_completion.py +++ b/litellm/tests/test_completion.py @@ -268,7 +268,7 @@ def test_completion_azure_gpt4_vision(): pytest.fail(f"Error occurred: {e}") -test_completion_azure_gpt4_vision() +# test_completion_azure_gpt4_vision() @pytest.mark.skip(reason="this test is flaky") @@ -990,9 +990,9 @@ def tester( print("azure ad token respoonse\n") print(response) litellm.input_callback = [] - except: + except Exception as e: litellm.input_callback = [] - pass + pytest.fail(f"An exception occurs - {str(e)}") # test_azure_openai_ad_token()