Skip to content

Commit

Permalink
Merge pull request #4936 from ConnectAI-E/feat-baidu
Browse files Browse the repository at this point in the history
feat: support baidu model
  • Loading branch information
lloydzhou authored Jul 9, 2024
2 parents 5295802 + f2a35f1 commit 47ea383
Show file tree
Hide file tree
Showing 15 changed files with 667 additions and 4 deletions.
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:
systemApiKey = serverConfig.baiduApiKey;
break;
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

0 comments on commit 47ea383

Please sign in to comment.