add azure support

app/api/chat-stream/route.ts

@@ -8,7 +8,8 @@ async function createStream(req: NextRequest) {
 
   const res = await requestOpenai(req);
 
-  const stream = new ReadableStream({
+  const stream = new ReadableStream(
+    {
       async start(controller) {
         function onParse(event: any) {
           if (event.type === "event") {
@@ -20,6 +21,12 @@ async function createStream(req: NextRequest) {
             }
             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);
@@ -34,7 +41,11 @@ async function createStream(req: NextRequest) {
           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 = {
-  return fetch(`${PROTOCOL}://${BASE_URL}/${openaiPath}`, {
-    headers: {
     "Content-Type": "application/json",
-      Authorization: `Bearer ${apiKey}`,
+  };
-    },
+  if (apiKey) {
+    headers.Authorization = `Bearer ${apiKey}`;
+  }
+  if (azureApiKey) {
+    headers["api-key"] = azureApiKey;
+  }
+
+  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(
         {