From ecf326d8906091d3ffe8ac6feebefda03cb4d1af Mon Sep 17 00:00:00 2001 From: Tongliang Liao Date: Tue, 26 Aug 2025 20:14:58 -0700 Subject: [PATCH] Support Azure OpenAI endpoints. Add AAD auth in addition to PSK. --- global_methods.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/global_methods.py b/global_methods.py index 7ada78d..4ff6eb5 100644 --- a/global_methods.py +++ b/global_methods.py @@ -23,6 +23,12 @@ def set_gemini_key(): def set_openai_key(): openai.api_key = os.environ['OPENAI_API_KEY'] + if openai.api_type == "azure" and openai.azure_ad_token_provider is None: + import azure.identity + openai.azure_ad_token_provider = azure.identity.get_bearer_token_provider( + azure.identity.DefaultAzureCredential(), + "https://cognitiveservices.azure.com/.default", + ) def run_json_trials(query, num_gen=1, num_tokens_request=1000,