Support Azure OpenAI API endpoint

src/khoj/processor/conversation/openai/utils.py

@@ -51,10 +51,17 @@ def completion_with_backoff(
     client_key = f"{openai_api_key}--{api_base_url}"
     client: openai.OpenAI | None = openai_clients.get(client_key)
     if not client:
-        client = openai.OpenAI(
-            api_key=openai_api_key,
-            base_url=api_base_url,
-        )
+        if api_base_url and "openai.azure.com" in api_base_url:
+            client = openai.AzureOpenAI(
+                api_key=openai_api_key,
+                azure_endpoint=api_base_url,
+                api_version="2024-10-21",
+            )
+        else:
+            client = openai.OpenAI(
+                api_key=openai_api_key,
+                base_url=api_base_url,
+            )
         openai_clients[client_key] = client
 
     formatted_messages = [{"role": message.role, "content": message.content} for message in messages]
@@ -158,14 +165,21 @@ def llm_thread(
 ):
     try:
         client_key = f"{openai_api_key}--{api_base_url}"
-        if client_key not in openai_clients:
-            client = openai.OpenAI(
-                api_key=openai_api_key,
-                base_url=api_base_url,
-            )
-            openai_clients[client_key] = client
-        else:
+        if client_key in openai_clients:
             client = openai_clients[client_key]
+        else:
+            if api_base_url and "openai.azure.com" in api_base_url:
+                client = openai.AzureOpenAI(
+                    api_key=openai_api_key,
+                    azure_endpoint=api_base_url,
+                    api_version="2024-10-21",
+                )
+            else:
+                client = openai.OpenAI(
+                    api_key=openai_api_key,
+                    base_url=api_base_url,
+                )
+            openai_clients[client_key] = client
 
         formatted_messages = [{"role": message.role, "content": message.content} for message in messages]