From 07915b68e78e972f2cf2150a70462f416174d3d8 Mon Sep 17 00:00:00 2001
From: Wataru Takahashi <19966304+wtberry@users.noreply.github.com>
Date: Fri, 27 Dec 2024 12:01:39 +0900
Subject: [PATCH] added support for azure openai api key

---
 promptwizard/glue/common/llm/llm_mgr.py | 34 +++++++++++++++++--------
 1 file changed, 23 insertions(+), 11 deletions(-)

diff --git a/promptwizard/glue/common/llm/llm_mgr.py b/promptwizard/glue/common/llm/llm_mgr.py
index c5cec9cd..e78d89a1 100644
--- a/promptwizard/glue/common/llm/llm_mgr.py
+++ b/promptwizard/glue/common/llm/llm_mgr.py
@@ -29,19 +29,31 @@ def call_api(messages):
         temperature=0.0,
         )
     else:
-        token_provider = get_bearer_token_provider(
-                AzureCliCredential(), "https://cognitiveservices.azure.com/.default"
+        if os.environ['AZURE_OPENAI_OPENAI_API_KEY']:
+            client = AzureOpenAI(
+                api_version=os.environ["OPENAI_API_VERSION"],
+                azure_endpoint=os.environ["AZURE_OPENAI_ENDPOINT"],
+                api_key=os.environ["AZURE_OPENAI_OPENAI_API_KEY"],
+                )
+            response = client.chat.completions.create(
+                model=os.environ["AZURE_OPENAI_DEPLOYMENT_NAME"],
+                messages=messages,
+                temperature=0.0,
             )
-        client = AzureOpenAI(
-            api_version=os.environ["OPENAI_API_VERSION"],
-            azure_endpoint=os.environ["AZURE_OPENAI_ENDPOINT"],
-            azure_ad_token_provider=token_provider
+        else:
+            token_provider = get_bearer_token_provider(
+                    AzureCliCredential(), "https://cognitiveservices.azure.com/.default"
+                )
+            client = AzureOpenAI(
+                api_version=os.environ["OPENAI_API_VERSION"],
+                azure_endpoint=os.environ["AZURE_OPENAI_ENDPOINT"],
+                azure_ad_token_provider=token_provider
+                )
+            response = client.chat.completions.create(
+                model=os.environ["AZURE_OPENAI_DEPLOYMENT_NAME"],
+                messages=messages,
+                temperature=0.0,
             )
-        response = client.chat.completions.create(
-            model=os.environ["AZURE_OPENAI_DEPLOYMENT_NAME"],
-            messages=messages,
-            temperature=0.0,
-        )
 
     prediction = response.choices[0].message.content
     return prediction