diff --git a/app/api/chat-stream/route.ts b/app/api/chat-stream/route.ts index e7bdfc5fb..3a2fa15f9 100644 --- a/app/api/chat-stream/route.ts +++ b/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; } diff --git a/app/api/common.ts b/app/api/common.ts index 842eeacaf..120511ef1 100644 --- a/app/api/common.ts +++ b/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, }); diff --git a/middleware.ts b/middleware.ts index 9338a2c6b..7515702d1 100644 --- a/middleware.ts +++ b/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( {