add azure support

This commit is contained in:
Zhe Li 2023-03-30 23:34:00 +08:00
parent 9eb77207fb
commit 5deedb26ba
3 changed files with 62 additions and 29 deletions

View File

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

View File

@ -8,14 +8,27 @@ const BASE_URL = process.env.BASE_URL ?? OPENAI_URL;
export async function requestOpenai(req: NextRequest) { export async function requestOpenai(req: NextRequest) {
const apiKey = req.headers.get("token"); const apiKey = req.headers.get("token");
const openaiPath = req.headers.get("path"); 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}`, { console.log("[Proxy] ", endpoint);
headers: {
"Content-Type": "application/json", return fetch(endpoint, {
Authorization: `Bearer ${apiKey}`, headers,
},
method: req.method, method: req.method,
body: req.body, body: req.body,
}); });

View File

@ -31,9 +31,18 @@ export function middleware(req: NextRequest) {
// inject api key // inject api key
if (!token) { if (!token) {
const apiKey = process.env.OPENAI_API_KEY; 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) { if (apiKey) {
console.log("[Auth] set system token"); console.log("[Auth] set system token");
req.headers.set("token", apiKey); 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 { } else {
return NextResponse.json( return NextResponse.json(
{ {