feat: close #3187 use CUSTOM_MODELS to control model list

This commit is contained in:
Yidadaa
2023-11-09 03:01:29 +08:00
parent fbc0236748
commit d93f05f511
11 changed files with 112 additions and 42 deletions

View File

@@ -1,10 +1,9 @@
import { NextRequest, NextResponse } from "next/server";
import { getServerSideConfig } from "../config/server";
import { DEFAULT_MODELS, OPENAI_BASE_URL } from "../constant";
import { collectModelTable, collectModels } from "../utils/model";
export const OPENAI_URL = "api.openai.com";
const DEFAULT_PROTOCOL = "https";
const PROTOCOL = process.env.PROTOCOL || DEFAULT_PROTOCOL;
const BASE_URL = process.env.BASE_URL || OPENAI_URL;
const DISABLE_GPT4 = !!process.env.DISABLE_GPT4;
const serverConfig = getServerSideConfig();
export async function requestOpenai(req: NextRequest) {
const controller = new AbortController();
@@ -14,10 +13,10 @@ export async function requestOpenai(req: NextRequest) {
"",
);
let baseUrl = BASE_URL;
let baseUrl = serverConfig.baseUrl ?? OPENAI_BASE_URL;
if (!baseUrl.startsWith("http")) {
baseUrl = `${PROTOCOL}://${baseUrl}`;
baseUrl = `https://${baseUrl}`;
}
if (baseUrl.endsWith("/")) {
@@ -26,10 +25,7 @@ export async function requestOpenai(req: NextRequest) {
console.log("[Proxy] ", openaiPath);
console.log("[Base Url]", baseUrl);
if (process.env.OPENAI_ORG_ID) {
console.log("[Org ID]", process.env.OPENAI_ORG_ID);
}
console.log("[Org ID]", serverConfig.openaiOrgId);
const timeoutId = setTimeout(
() => {
@@ -58,18 +54,23 @@ export async function requestOpenai(req: NextRequest) {
};
// #1815 try to refuse gpt4 request
if (DISABLE_GPT4 && req.body) {
if (serverConfig.customModels && req.body) {
try {
const modelTable = collectModelTable(
DEFAULT_MODELS,
serverConfig.customModels,
);
const clonedBody = await req.text();
fetchOptions.body = clonedBody;
const jsonBody = JSON.parse(clonedBody);
const jsonBody = JSON.parse(clonedBody) as { model?: string };
if ((jsonBody?.model ?? "").includes("gpt-4")) {
// not undefined and is false
if (modelTable[jsonBody?.model ?? ""] === false) {
return NextResponse.json(
{
error: true,
message: "you are not allowed to use gpt-4 model",
message: `you are not allowed to use ${jsonBody?.model} model`,
},
{
status: 403,