perf: bill framwork

src/service/events/generateQA.ts

@@ -9,6 +9,7 @@ import { SplitDataSchema } from '@/types/mongoSchema';
 import { modelServiceToolMap } from '../utils/chat';
 import { ChatRoleEnum } from '@/constants/chat';
 import { getErrText } from '@/utils/tools';
+import { BillTypeEnum } from '@/constants/user';
 
 export async function generateQA(next = false): Promise<any> {
   if (process.env.queueTask !== '1') {
@@ -98,7 +99,7 @@ A2:
             pushSplitDataBill({
               isPay: !userOpenAiKey && result.length > 0,
               userId: dataItem.userId,
-              type: 'QA',
+              type: BillTypeEnum.QA,
               textLen: responseMessages.map((item) => item.value).join('').length,
               totalTokens
             });

src/service/events/pushBill.ts

@@ -8,7 +8,8 @@ export const pushChatBill = async ({
   userId,
   chatId,
   textLen,
-  tokens
+  tokens,
+  type
 }: {
   isPay: boolean;
   chatModel: ChatModelType;
@@ -16,6 +17,7 @@ export const pushChatBill = async ({
   chatId?: '' | string;
   textLen: number;
   tokens: number;
+  type: BillTypeEnum.chat | BillTypeEnum.openapiChat;
 }) => {
   console.log(`chat generate success. text len: ${textLen}. token len: ${tokens}. pay:${isPay}`);
   if (!isPay) return;
@@ -33,7 +35,7 @@ export const pushChatBill = async ({
       // 插入 Bill 记录
       const res = await Bill.create({
         userId,
-        type: 'chat',
+        type,
         modelName: chatModel,
         chatId: chatId ? chatId : undefined,
         textLen,
@@ -83,7 +85,7 @@ export const pushSplitDataBill = async ({
   userId: string;
   totalTokens: number;
   textLen: number;
-  type: `${BillTypeEnum}`;
+  type: BillTypeEnum.QA;
 }) => {
   console.log(
     `splitData generate success. text len: ${textLen}. token len: ${totalTokens}. pay:${isPay}`

src/service/models/bill.ts

@@ -1,5 +1,5 @@
 import { Schema, model, models, Model } from 'mongoose';
-import { ChatModelMap } from '@/constants/model';
+import { ChatModelMap, embeddingModel } from '@/constants/model';
 import { BillSchema as BillType } from '@/types/mongoSchema';
 import { BillTypeMap } from '@/constants/user';
 
@@ -16,7 +16,7 @@ const BillSchema = new Schema({
   },
   modelName: {
     type: String,
-    enum: [...Object.keys(ChatModelMap), 'text-embedding-ada-002'],
+    enum: [...Object.keys(ChatModelMap), embeddingModel],
     required: true
   },
   chatId: {

src/service/plugins/useSearch.ts

@@ -1,35 +1,34 @@
 import axios from 'axios';
 
-{/*Bing 搜索*/}
+{
+  /*Bing 搜索*/
+}
 const BingSearch = async (wait_val: string) => {
-  const response = await axios.post("newbing中转服务器", {
-    prompt: wait_val,
+  const response = await axios.post('newbing中转服务器', {
+    prompt: wait_val
   });
   const result = response.data.result;
   return result;
 };
 
-
-{/*google 搜索*/}
+{
+  /*google 搜索*/
+}
 const GoogleSearch = async (wait_val: string) => {
-
   const response = await axios.get('https://www.googleapis.com/customsearch/v1', {
     params: {
-      key: google_api,
+      key: process.env.GOOGLE_KEY,
       q: wait_val,
-      cx: searchEngineId,
+      cx: process.env.searchEngineId,
       start: 1,
       num: 3,
-      dateRestrict: 'm[1]', //搜索结果限定为一个月内
-    },
+      dateRestrict: 'm[1]' //搜索结果限定为一个月内
+    }
   });
   const results = response.data.items;
   if (results !== null) {
     const result = results.map((item: { snippet: string }) => item.snippet).join('\n');
     return result;
   }
-  }
-export {
-  BingSearch,
-  GoogleSearch
 };
+export { BingSearch, GoogleSearch };

src/service/utils/auth.ts

@@ -14,13 +14,11 @@ import { hashPassword } from '@/service/utils/tools';
 /* uniform auth user */
 export const authUser = async ({
   req,
-  userId = '',
   authToken = false,
   authOpenApi = false,
   authRoot = false
 }: {
   req: NextApiRequest;
-  userId?: string;
   authToken?: boolean;
   authOpenApi?: boolean;
   authRoot?: boolean;
@@ -68,17 +66,18 @@ export const authUser = async ({
       return Promise.reject(error);
     }
   };
-  const parseRootKey = async (rootKey?: string) => {
-    if (!rootKey || !process.env.ROOT_KEY || rootKey !== process.env.ROOT_KEY) {
+  const parseRootKey = async (rootKey?: string, userId?: string) => {
+    if (!rootKey || !userId || !process.env.ROOT_KEY || rootKey !== process.env.ROOT_KEY) {
       return Promise.reject(ERROR_ENUM.unAuthorization);
     }
     return userId;
   };
 
-  const { cookie, apikey, rootkey } = (req.headers || {}) as {
+  const { cookie, apikey, rootkey, userid } = (req.headers || {}) as {
     cookie?: string;
     apikey?: string;
     rootkey?: string;
+    userid?: string;
   };
 
   let uid = '';
@@ -88,13 +87,13 @@ export const authUser = async ({
   } else if (authOpenApi) {
     uid = await parseOpenApiKey(apikey);
   } else if (authRoot) {
-    uid = await parseRootKey(rootkey);
+    uid = await parseRootKey(rootkey, userid);
   } else if (cookie) {
     uid = await parseCookie(cookie);
   } else if (apikey) {
     uid = await parseOpenApiKey(apikey);
   } else if (rootkey) {
-    uid = await parseRootKey(rootkey);
+    uid = await parseRootKey(rootkey, userid);
   } else {
     return Promise.reject(ERROR_ENUM.unAuthorization);
   }