Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: support baidu model #4936

Merged
merged 11 commits into from
Jul 9, 2024
Merged
2 changes: 1 addition & 1 deletion .gitignore
Original file line number Diff line number Diff line change
Expand Up @@ -43,4 +43,4 @@ dev
.env

*.key
*.key.pub
*.key.pub
12 changes: 12 additions & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -213,6 +213,18 @@ anthropic claude Api version.

anthropic claude Api Url.

### `BAIDU_API_KEY` (optional)

Baidu Api Key.

### `BAIDU_SECRET_KEY` (optional)

Baidu Secret Key.

### `BAIDU_URL` (optional)

Baidu Api Url.

### `HIDE_USER_API_KEY` (optional)

> Default: Empty
Expand Down
12 changes: 12 additions & 0 deletions README_CN.md
Original file line number Diff line number Diff line change
Expand Up @@ -127,6 +127,18 @@ anthropic claude Api version.

anthropic claude Api Url.

### `BAIDU_API_KEY` (可选)

Baidu Api Key.

### `BAIDU_SECRET_KEY` (可选)

Baidu Secret Key.

### `BAIDU_URL` (可选)

Baidu Api Url.

### `HIDE_USER_API_KEY` (可选)

如果你不想让用户自行填入 API Key,将此环境变量设置为 1 即可。
Expand Down
3 changes: 3 additions & 0 deletions app/api/auth.ts
Original file line number Diff line number Diff line change
Expand Up @@ -73,6 +73,9 @@ export function auth(req: NextRequest, modelProvider: ModelProvider) {
case ModelProvider.Claude:
systemApiKey = serverConfig.anthropicApiKey;
break;
case ModelProvider.Ernie:
lloydzhou marked this conversation as resolved.
Show resolved Hide resolved
systemApiKey = serverConfig.baiduApiKey;
break;
lloydzhou marked this conversation as resolved.
Show resolved Hide resolved
case ModelProvider.GPT:
default:
if (req.nextUrl.pathname.includes("azure/deployments")) {
Expand Down
169 changes: 169 additions & 0 deletions app/api/baidu/[...path]/route.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,169 @@
import { getServerSideConfig } from "@/app/config/server";
import {
BAIDU_BASE_URL,
ApiPath,
ModelProvider,
BAIDU_OATUH_URL,
ServiceProvider,
} from "@/app/constant";
import { prettyObject } from "@/app/utils/format";
import { NextRequest, NextResponse } from "next/server";
import { auth } from "@/app/api/auth";
import { isModelAvailableInServer } from "@/app/utils/model";
import { getAccessToken } from "@/app/utils/baidu";

const serverConfig = getServerSideConfig();

async function handle(
req: NextRequest,
{ params }: { params: { path: string[] } },
) {
console.log("[Baidu Route] params ", params);

if (req.method === "OPTIONS") {
return NextResponse.json({ body: "OK" }, { status: 200 });
}

const authResult = auth(req, ModelProvider.Ernie);
if (authResult.error) {
return NextResponse.json(authResult, {
status: 401,
});
}

if (!serverConfig.baiduApiKey || !serverConfig.baiduSecretKey) {
return NextResponse.json(
{
error: true,
message: `missing BAIDU_API_KEY or BAIDU_SECRET_KEY in server env vars`,
},
{
status: 401,
},
);
}

try {
const response = await request(req);
return response;
} catch (e) {
console.error("[Baidu] ", e);
return NextResponse.json(prettyObject(e));
}
}

export const GET = handle;
export const POST = handle;

export const runtime = "edge";
export const preferredRegion = [
"arn1",
"bom1",
"cdg1",
"cle1",
"cpt1",
"dub1",
"fra1",
"gru1",
"hnd1",
"iad1",
"icn1",
"kix1",
"lhr1",
"pdx1",
"sfo1",
"sin1",
"syd1",
];

async function request(req: NextRequest) {
const controller = new AbortController();

let path = `${req.nextUrl.pathname}`.replaceAll(ApiPath.Baidu, "");

let baseUrl = serverConfig.baiduUrl || BAIDU_BASE_URL;

if (!baseUrl.startsWith("http")) {
baseUrl = `https://${baseUrl}`;
}

if (baseUrl.endsWith("/")) {
baseUrl = baseUrl.slice(0, -1);
}

console.log("[Proxy] ", path);
console.log("[Base Url]", baseUrl);

const timeoutId = setTimeout(
() => {
controller.abort();
},
10 * 60 * 1000,
);

const { access_token } = await getAccessToken(
serverConfig.baiduApiKey as string,
serverConfig.baiduSecretKey as string,
);
const fetchUrl = `${baseUrl}${path}?access_token=${access_token}`;

const fetchOptions: RequestInit = {
headers: {
"Content-Type": "application/json",
},
method: req.method,
body: req.body,
redirect: "manual",
// @ts-ignore
duplex: "half",
signal: controller.signal,
};

// #1815 try to refuse some request to some models
if (serverConfig.customModels && req.body) {
try {
const clonedBody = await req.text();
fetchOptions.body = clonedBody;

const jsonBody = JSON.parse(clonedBody) as { model?: string };

// not undefined and is false
if (
isModelAvailableInServer(
serverConfig.customModels,
jsonBody?.model as string,
ServiceProvider.Baidu as string,
)
) {
return NextResponse.json(
{
error: true,
message: `you are not allowed to use ${jsonBody?.model} model`,
},
{
status: 403,
},
);
}
} catch (e) {
console.error(`[Baidu] filter`, e);
}
}
try {
const res = await fetch(fetchUrl, fetchOptions);

// to prevent browser prompt for credentials
const newHeaders = new Headers(res.headers);
newHeaders.delete("www-authenticate");
// to disable nginx buffering
newHeaders.set("X-Accel-Buffering", "no");

return new Response(res.body, {
status: res.status,
statusText: res.statusText,
headers: newHeaders,
});
} finally {
clearTimeout(timeoutId);
}
}
6 changes: 4 additions & 2 deletions app/api/common.ts
Original file line number Diff line number Diff line change
Expand Up @@ -70,7 +70,7 @@ export async function requestOpenai(req: NextRequest) {
// Forward compatibility:
// if display_name(deployment_name) not set, and '{deploy-id}' in AZURE_URL
// then using default '{deploy-id}'
if (serverConfig.customModels) {
if (serverConfig.customModels && serverConfig.azureUrl) {
const modelName = path.split("/")[1];
let realDeployName = "";
serverConfig.customModels
Expand All @@ -80,7 +80,9 @@ export async function requestOpenai(req: NextRequest) {
const [fullName, displayName] = m.split("=");
const [_, providerName] = fullName.split("@");
if (providerName === "azure" && !displayName) {
const [_, deployId] = serverConfig.azureUrl.split("deployments/");
const [_, deployId] = (serverConfig?.azureUrl ?? "").split(
"deployments/",
);
if (deployId) {
realDeployName = deployId;
}
Expand Down
7 changes: 7 additions & 0 deletions app/client/api.ts
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,8 @@ import { ChatMessage, ModelType, useAccessStore, useChatStore } from "../store";
import { ChatGPTApi } from "./platforms/openai";
import { GeminiProApi } from "./platforms/google";
import { ClaudeApi } from "./platforms/anthropic";
import { ErnieApi } from "./platforms/baidu";

export const ROLES = ["system", "user", "assistant"] as const;
export type MessageRole = (typeof ROLES)[number];

Expand Down Expand Up @@ -104,6 +106,9 @@ export class ClientApi {
case ModelProvider.Claude:
this.llm = new ClaudeApi();
break;
case ModelProvider.Ernie:
this.llm = new ErnieApi();
break;
default:
this.llm = new ChatGPTApi();
}
Expand Down Expand Up @@ -220,6 +225,8 @@ export function getClientApi(provider: ServiceProvider): ClientApi {
return new ClientApi(ModelProvider.GeminiPro);
case ServiceProvider.Anthropic:
return new ClientApi(ModelProvider.Claude);
case ServiceProvider.Baidu:
return new ClientApi(ModelProvider.Ernie);
default:
return new ClientApi(ModelProvider.GPT);
}
Expand Down
Loading
Loading