From c879905307c8910a1e25543170cc623d17456a46 Mon Sep 17 00:00:00 2001
From: archer <545436317@qq.com>
Date: Thu, 4 May 2023 12:18:07 +0800
Subject: [PATCH] feat: openapi auth

---
 src/pages/api/openapi/chat/chat.ts   | 17 ++++++++++++---
 src/pages/api/openapi/chat/lafGpt.ts | 21 ++++++++++++------
 src/service/errorCode.ts             |  2 +-
 src/service/utils/auth.ts            | 32 +++++++++++-----------------
 4 files changed, 41 insertions(+), 31 deletions(-)

diff --git a/src/pages/api/openapi/chat/chat.ts b/src/pages/api/openapi/chat/chat.ts
index 58ac64956..0c724a4a1 100644
--- a/src/pages/api/openapi/chat/chat.ts
+++ b/src/pages/api/openapi/chat/chat.ts
@@ -1,6 +1,6 @@
 import type { NextApiRequest, NextApiResponse } from 'next';
 import { connectToDatabase } from '@/service/mongo';
-import { authOpenApiKey, authModel } from '@/service/utils/auth';
+import { authOpenApiKey, authModel, getApiKey } from '@/service/utils/auth';
 import { modelServiceToolMap, resStreamResponse } from '@/service/utils/chat';
 import { ChatItemSimpleType } from '@/types/chat';
 import { jsonRes } from '@/service/response';
@@ -28,10 +28,12 @@ export default async function handler(req: NextApiRequest, res: NextApiResponse)
 
   try {
     const {
+      chatId,
       prompts,
       modelId,
       isStream = true
     } = req.body as {
+      chatId?: string;
       prompts: ChatItemSimpleType[];
       modelId: string;
       isStream: boolean;
@@ -51,13 +53,20 @@ export default async function handler(req: NextApiRequest, res: NextApiResponse)
     let startTime = Date.now();
 
     /* 凭证校验 */
-    const { apiKey, userId } = await authOpenApiKey(req);
+    const { userId } = await authOpenApiKey(req);
 
     const { model } = await authModel({
       userId,
       modelId
     });
 
+    /* get api key */
+    const { systemAuthKey: apiKey } = await getApiKey({
+      model: model.chat.chatModel,
+      userId,
+      mustPay: true
+    });
+
     const modelConstantsData = ChatModelMap[model.chat.chatModel];
 
     // 使用了知识库搜索
@@ -98,7 +107,9 @@ export default async function handler(req: NextApiRequest, res: NextApiResponse)
         apiKey,
         temperature: +temperature,
         messages: prompts,
-        stream: isStream
+        stream: isStream,
+        res,
+        chatId
       });
 
     console.log('api response time:', `${(Date.now() - startTime) / 1000}s`);
diff --git a/src/pages/api/openapi/chat/lafGpt.ts b/src/pages/api/openapi/chat/lafGpt.ts
index c0e3542ea..15a37f82c 100644
--- a/src/pages/api/openapi/chat/lafGpt.ts
+++ b/src/pages/api/openapi/chat/lafGpt.ts
@@ -1,6 +1,6 @@
 import type { NextApiRequest, NextApiResponse } from 'next';
-import { connectToDatabase, Model } from '@/service/mongo';
-import { authOpenApiKey } from '@/service/utils/auth';
+import { connectToDatabase } from '@/service/mongo';
+import { authOpenApiKey, authModel, getApiKey } from '@/service/utils/auth';
 import { resStreamResponse, modelServiceToolMap } from '@/service/utils/chat';
 import { ChatItemSimpleType } from '@/types/chat';
 import { jsonRes } from '@/service/response';
@@ -45,13 +45,20 @@ export default async function handler(req: NextApiRequest, res: NextApiResponse)
     let startTime = Date.now();
 
     /* 凭证校验 */
-    const { apiKey, userId } = await authOpenApiKey(req);
+    const { userId } = await authOpenApiKey(req);
 
     /* 查找数据库里的模型信息 */
-    const model = await Model.findById(modelId);
-    if (!model) {
-      throw new Error('找不到模型');
-    }
+    const { model } = await authModel({
+      userId,
+      modelId
+    });
+
+    /* get api key */
+    const { systemAuthKey: apiKey } = await getApiKey({
+      model: model.chat.chatModel,
+      userId,
+      mustPay: true
+    });
 
     const modelConstantsData = ChatModelMap[model.chat.chatModel];
 
diff --git a/src/service/errorCode.ts b/src/service/errorCode.ts
index 0879c9c3e..52031b23c 100644
--- a/src/service/errorCode.ts
+++ b/src/service/errorCode.ts
@@ -56,7 +56,7 @@ export const ERROR_RESPONSE: Record<
     data: null
   },
   [ERROR_ENUM.insufficientQuota]: {
-    code: 403,
+    code: 510,
     statusText: ERROR_ENUM.insufficientQuota,
     message: '账号余额不足',
     data: null
diff --git a/src/service/utils/auth.ts b/src/service/utils/auth.ts
index f2b426134..2211c5327 100644
--- a/src/service/utils/auth.ts
+++ b/src/service/utils/auth.ts
@@ -29,13 +29,18 @@ export const authToken = (token?: string): Promise<string> => {
 };
 
 /* 获取 api 请求的 key */
-export const getApiKey = async ({ model, userId }: { model: ChatModelType; userId: string }) => {
+export const getApiKey = async ({
+  model,
+  userId,
+  mustPay = false
+}: {
+  model: ChatModelType;
+  userId: string;
+  mustPay?: boolean;
+}) => {
   const user = await User.findById(userId);
   if (!user) {
-    return Promise.reject({
-      code: 501,
-      message: '找不到用户'
-    });
+    return Promise.reject(ERROR_ENUM.unAuthorization);
   }
 
   const keyMap = {
@@ -58,7 +63,7 @@ export const getApiKey = async ({ model, userId }: { model: ChatModelType; userI
   };
 
   // 有自己的key
-  if (keyMap[model].userOpenAiKey) {
+  if (!mustPay && keyMap[model].userOpenAiKey) {
     return {
       user,
       userOpenAiKey: keyMap[model].userOpenAiKey,
@@ -68,10 +73,7 @@ export const getApiKey = async ({ model, userId }: { model: ChatModelType; userI
 
   // 平台账号余额校验
   if (formatPrice(user.balance) <= 0) {
-    return Promise.reject({
-      code: 501,
-      message: '账号余额不足'
-    });
+    return Promise.reject(ERROR_ENUM.unAuthorization);
   }
 
   return {
@@ -192,22 +194,12 @@ export const authOpenApiKey = async (req: NextApiRequest) => {
     }
     const userId = String(openApi.userId);
 
-    // 余额校验
-    const user = await User.findById(userId);
-    if (!user) {
-      return Promise.reject(ERROR_ENUM.unAuthorization);
-    }
-    if (formatPrice(user.balance) <= 0) {
-      return Promise.reject(ERROR_ENUM.insufficientQuota);
-    }
-
     // 更新使用的时间
     await OpenApi.findByIdAndUpdate(openApi._id, {
       lastUsedTime: new Date()
     });
 
     return {
-      apiKey: process.env.OPENAIKEY as string,
       userId
     };
   } catch (error) {