Merge pull request #3625 from fredliang44/main

fix: fix auth key selection when using different model
This commit is contained in:
Fred Liang 2023-12-25 05:13:38 +08:00 committed by GitHub
commit 771cc9e583
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 38 additions and 13 deletions

View File

@ -1,7 +1,7 @@
import { NextRequest } from "next/server"; import { NextRequest } from "next/server";
import { getServerSideConfig } from "../config/server"; import { getServerSideConfig } from "../config/server";
import md5 from "spark-md5"; import md5 from "spark-md5";
import { ACCESS_CODE_PREFIX } from "../constant"; import { ACCESS_CODE_PREFIX, ModelProvider } from "../constant";
function getIP(req: NextRequest) { function getIP(req: NextRequest) {
let ip = req.ip ?? req.headers.get("x-real-ip"); let ip = req.ip ?? req.headers.get("x-real-ip");
@ -24,7 +24,7 @@ function parseApiKey(bearToken: string) {
}; };
} }
export function auth(req: NextRequest) { export function auth(req: NextRequest, modelProvider: ModelProvider) {
const authToken = req.headers.get("Authorization") ?? ""; const authToken = req.headers.get("Authorization") ?? "";
// check if it is openai api key or user token // check if it is openai api key or user token
@ -56,12 +56,19 @@ export function auth(req: NextRequest) {
// if user does not provide an api key, inject system api key // if user does not provide an api key, inject system api key
if (!apiKey) { if (!apiKey) {
const serverConfig = getServerSideConfig(); const serverConfig = getServerSideConfig();
const systemApiKey = serverConfig.isAzure
? serverConfig.azureApiKey
: serverConfig.isGoogle
? serverConfig.googleApiKey
: serverConfig.apiKey;
// const systemApiKey = serverConfig.isAzure
// ? serverConfig.azureApiKey
// : serverConfig.isGoogle
// ? serverConfig.googleApiKey
// : serverConfig.apiKey;
const systemApiKey =
modelProvider === ModelProvider.GeminiPro
? serverConfig.googleApiKey
: serverConfig.isAzure
? serverConfig.azureApiKey
: serverConfig.apiKey;
if (systemApiKey) { if (systemApiKey) {
console.log("[Auth] use system api key"); console.log("[Auth] use system api key");
req.headers.set("Authorization", `Bearer ${systemApiKey}`); req.headers.set("Authorization", `Bearer ${systemApiKey}`);

View File

@ -1,7 +1,7 @@
import { NextRequest, NextResponse } from "next/server"; import { NextRequest, NextResponse } from "next/server";
import { auth } from "../../auth"; import { auth } from "../../auth";
import { getServerSideConfig } from "@/app/config/server"; import { getServerSideConfig } from "@/app/config/server";
import { GEMINI_BASE_URL, Google } from "@/app/constant"; import { GEMINI_BASE_URL, Google, ModelProvider } from "@/app/constant";
async function handle( async function handle(
req: NextRequest, req: NextRequest,
@ -39,7 +39,7 @@ async function handle(
10 * 60 * 1000, 10 * 60 * 1000,
); );
const authResult = auth(req); const authResult = auth(req, ModelProvider.GeminiPro);
if (authResult.error) { if (authResult.error) {
return NextResponse.json(authResult, { return NextResponse.json(authResult, {
status: 401, status: 401,
@ -50,6 +50,7 @@ async function handle(
const token = bearToken.trim().replaceAll("Bearer ", "").trim(); const token = bearToken.trim().replaceAll("Bearer ", "").trim();
const key = token ? token : serverConfig.googleApiKey; const key = token ? token : serverConfig.googleApiKey;
if (!key) { if (!key) {
return NextResponse.json( return NextResponse.json(
{ {
@ -63,7 +64,6 @@ async function handle(
} }
const fetchUrl = `${baseUrl}/${path}?key=${key}`; const fetchUrl = `${baseUrl}/${path}?key=${key}`;
const fetchOptions: RequestInit = { const fetchOptions: RequestInit = {
headers: { headers: {
"Content-Type": "application/json", "Content-Type": "application/json",

View File

@ -1,6 +1,6 @@
import { type OpenAIListModelResponse } from "@/app/client/platforms/openai"; import { type OpenAIListModelResponse } from "@/app/client/platforms/openai";
import { getServerSideConfig } from "@/app/config/server"; import { getServerSideConfig } from "@/app/config/server";
import { OpenaiPath } from "@/app/constant"; import { ModelProvider, OpenaiPath } from "@/app/constant";
import { prettyObject } from "@/app/utils/format"; import { prettyObject } from "@/app/utils/format";
import { NextRequest, NextResponse } from "next/server"; import { NextRequest, NextResponse } from "next/server";
import { auth } from "../../auth"; import { auth } from "../../auth";
@ -45,7 +45,7 @@ async function handle(
); );
} }
const authResult = auth(req); const authResult = auth(req, ModelProvider.GPT);
if (authResult.error) { if (authResult.error) {
return NextResponse.json(authResult, { return NextResponse.json(authResult, {
status: 401, status: 401,
@ -75,4 +75,22 @@ export const GET = handle;
export const POST = handle; export const POST = handle;
export const runtime = "edge"; export const runtime = "edge";
export const preferredRegion = ['arn1', 'bom1', 'cdg1', 'cle1', 'cpt1', 'dub1', 'fra1', 'gru1', 'hnd1', 'iad1', 'icn1', 'kix1', 'lhr1', 'pdx1', 'sfo1', 'sin1', 'syd1']; export const preferredRegion = [
"arn1",
"bom1",
"cdg1",
"cle1",
"cpt1",
"dub1",
"fra1",
"gru1",
"hnd1",
"iad1",
"icn1",
"kix1",
"lhr1",
"pdx1",
"sfo1",
"sin1",
"syd1",
];