From e3ae67f8e238b8a00e997599ffb206d953261fd5 Mon Sep 17 00:00:00 2001
From: Hugues Chocart <chocart.hugues@icloud.com>
Date: Tue, 16 Apr 2024 13:18:41 -0300
Subject: [PATCH] feat: add azure openai support (#231)

---
 packages/backend/src/utils/openai.ts     | 26 +++++++++++++++++++++---
 packages/backend/src/utils/playground.ts |  5 ++---
 2 files changed, 25 insertions(+), 6 deletions(-)

diff --git a/packages/backend/src/utils/openai.ts b/packages/backend/src/utils/openai.ts
index 8a639a69..eae59261 100644
--- a/packages/backend/src/utils/openai.ts
+++ b/packages/backend/src/utils/openai.ts
@@ -1,6 +1,26 @@
-import OpenAI from "openai"
 import { monitorOpenAI } from "lunary/openai"
+import OpenAI from "openai"
+
+export function getOpenAIParams() {
+  if (process.env.OPENAI_API_KEY) {
+    return {
+      apiKey: process.env.OPENAI_API_KEY,
+    }
+  } else if (process.env.AZURE_OPENAI_API_KEY) {
+    const apiKey = process.env.AZURE_OPENAI_API_KEY
+    const model = process.env.AZURE_OPENAI_DEPLOYMENT_ID
+    const resource = process.env.AZURE_OPENAI_RESOURCE_NAME
 
-const openai = process.env.OPENAI_API_KEY ? monitorOpenAI(new OpenAI()) : null
+    return {
+      apiKey,
+      baseURL: `https://${resource}.openai.azure.com/openai/deployments/${model}`,
+      defaultQuery: { "api-version": "2024-02-01" },
+      defaultHeaders: { "api-key": apiKey },
+    }
+  } else {
+    return null
+  }
+}
+const clientParams = getOpenAIParams()
 
-export default openai
+export default clientParams ? monitorOpenAI(new OpenAI(clientParams)) : null
diff --git a/packages/backend/src/utils/playground.ts b/packages/backend/src/utils/playground.ts
index 18266c29..03c96afc 100644
--- a/packages/backend/src/utils/playground.ts
+++ b/packages/backend/src/utils/playground.ts
@@ -2,6 +2,7 @@ import { clearUndefined } from "./ingest"
 import OpenAI from "openai"
 import { MODELS } from "shared"
 import { ReadableStream } from "stream/web"
+import { getOpenAIParams } from "./openai"
 
 function convertInputToOpenAIMessages(input: any[]) {
   return input.map(({ role, content, text, functionCall, toolCalls, name }) => {
@@ -228,9 +229,7 @@ export async function runAImodel(
       break
 
     case "openai":
-      clientParams = {
-        apiKey: process.env.OPENAI_API_KEY,
-      }
+      clientParams = getOpenAIParams()
       break
     case "mistral":
       clientParams = {