Support AzureOpenAI client

Fixes #3

action.yml

@@ -43,6 +43,14 @@ inputs:
     description: 'Overwrite the PR description if it already exists. This also means that the action will run on every PR update'
     required: false
     default: 'false'
+  azure_endpoint:
+    description: 'Azure endpoint for OpenAI, if you are using Azure. Setting this will make the action use Azure OpenAI.'
+    required: false
+    default: ''
+  azure_openai_api_version:
+    description: 'OpenAI API version, if you are using Azure.'
+    required: false
+    default: ''
 
 runs:
   using: 'docker'

autofill_description.py

@@ -176,7 +176,18 @@ def main():
     if len(completion_prompt) > max_allowed_characters:
         completion_prompt = completion_prompt[:max_allowed_characters]
 
-    openai_client = openai.OpenAI(api_key=openai_api_key)
+    azure_endpoint = os.environ.get("INPUT_AZURE_ENDPOINT", "")
+    azure_openai_api_version = os.environ.get("INPUT_AZURE_OPENAI_API_VERSION", "")
+
+    openai_client = (
+        openai.AzureOpenAI(
+            api_key=openai_api_key,
+            azure_endpoint=azure_endpoint,
+            api_version=azure_openai_api_version,
+        )
+        if azure_endpoint
+        else openai.OpenAI(api_key=openai_api_key)
+    )
     openai_response = openai_client.chat.completions.create(
         model=open_ai_model,
         messages=[