add azure support

app/api/chat-stream/route.ts

@@ -8,33 +8,44 @@ async function createStream(req: NextRequest) {
 
   const res = await requestOpenai(req);
 
-  const stream = new ReadableStream({
-    async start(controller) {
-      function onParse(event: any) {
-        if (event.type === "event") {
-          const data = event.data;
-          // https://beta.openai.com/docs/api-reference/completions/create#completions/create-stream
-          if (data === "[DONE]") {
-            controller.close();
-            return;
-          }
-          try {
-            const json = JSON.parse(data);
-            const text = json.choices[0].delta.content;
-            const queue = encoder.encode(text);
-            controller.enqueue(queue);
-          } catch (e) {
-            controller.error(e);
+  const stream = new ReadableStream(
+    {
+      async start(controller) {
+        function onParse(event: any) {
+          if (event.type === "event") {
+            const data = event.data;
+            // https://beta.openai.com/docs/api-reference/completions/create#completions/create-stream
+            if (data === "[DONE]") {
+              controller.close();
+              return;
+            }
+            try {
+              const json = JSON.parse(data);
+
+              if (json.choices[0]["finish_reason"] === "stop") {
+                controller.close();
+                return;
+              }
+
+              const text = json.choices[0].delta.content;
+              const queue = encoder.encode(text);
+              controller.enqueue(queue);
+            } catch (e) {
+              controller.error(e);
+            }
           }
         }
-      }
 
-      const parser = createParser(onParse);
-      for await (const chunk of res.body as any) {
-        parser.feed(decoder.decode(chunk));
-      }
+        const parser = createParser(onParse);
+        for await (const chunk of res.body as any) {
+          parser.feed(decoder.decode(chunk));
+        }
+      },
     },
-  });
+    {
+      highWaterMark: 0,
+    },
+  );
   return stream;
 }
 

app/api/common.ts

@@ -8,14 +8,27 @@ const BASE_URL = process.env.BASE_URL ?? OPENAI_URL;
 export async function requestOpenai(req: NextRequest) {
   const apiKey = req.headers.get("token");
   const openaiPath = req.headers.get("path");
+  const openaiEndpoint = `${PROTOCOL}://${BASE_URL}/${openaiPath}`;
+  const azureApiKey = req.headers.get("azure-api-key");
+  const azureAccount = req.headers.get("azure-account");
+  const azureModel = req.headers.get("azure-model");
+  const azureEndpoint = `https://${azureAccount}.openai.azure.com/openai/deployments/${azureModel}/chat/completions?api-version=2023-03-15-preview`;
 
-  console.log("[Proxy] ", openaiPath);
+  const endpoint = apiKey ? openaiEndpoint : azureEndpoint;
+  const headers: HeadersInit = {
+    "Content-Type": "application/json",
+  };
+  if (apiKey) {
+    headers.Authorization = `Bearer ${apiKey}`;
+  }
+  if (azureApiKey) {
+    headers["api-key"] = azureApiKey;
+  }
 
-  return fetch(`${PROTOCOL}://${BASE_URL}/${openaiPath}`, {
-    headers: {
-      "Content-Type": "application/json",
-      Authorization: `Bearer ${apiKey}`,
-    },
+  console.log("[Proxy] ", endpoint);
+
+  return fetch(endpoint, {
+    headers,
     method: req.method,
     body: req.body,
   });

middleware.ts

@@ -31,9 +31,18 @@ export function middleware(req: NextRequest) {
   // inject api key
   if (!token) {
     const apiKey = process.env.OPENAI_API_KEY;
+    const azureApiKey = process.env.AZURE_API_KEY;
+    const azureAccount = process.env.AZURE_ACCOUNT;
+    const azureModel = process.env.AZURE_MODEL;
+
     if (apiKey) {
       console.log("[Auth] set system token");
       req.headers.set("token", apiKey);
+    } else if (azureApiKey && azureAccount && azureModel) {
+      console.log("[Auth] set system azure token");
+      req.headers.set("azure-api-key", azureApiKey);
+      req.headers.set("azure-account", azureAccount);
+      req.headers.set("azure-model", azureModel);
     } else {
       return NextResponse.json(
         {