fix(main.py): read azure ad token from optional params extra body

BerriAI · Jan 19, 2024 · e0aaa94 · e0aaa94
1 parent a8ba5df
commit e0aaa94
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 6 deletions.
diff --git 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
@@ -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()