feat: model related kb

src/service/models/model.ts

@@ -31,10 +31,10 @@ const ModelSchema = new Schema({
     default: () => new Date()
   },
   chat: {
-    useKb: {
-      // use knowledge base to search
-      type: Boolean,
-      default: false
+    relatedKbs: {
+      type: [Schema.Types.ObjectId],
+      ref: 'kb',
+      default: []
     },
     searchMode: {
       // knowledge base search mode
@@ -79,33 +79,6 @@ const ModelSchema = new Schema({
       type: Number,
       default: 0
     }
-  },
-  security: {
-    type: {
-      domain: {
-        type: [String],
-        default: ['*']
-      },
-      contextMaxLen: {
-        type: Number,
-        default: 20
-      },
-      contentMaxLen: {
-        type: Number,
-        default: 4000
-      },
-      expiredTime: {
-        type: Number,
-        default: 1,
-        set: (val: number) => val * (60 * 60 * 1000)
-      },
-      maxLoadAmount: {
-        // 负数代表不限制
-        type: Number,
-        default: -1
-      }
-    },
-    default: {}
   }
 });
 

src/service/plugins/searchKb.ts

@@ -48,7 +48,7 @@ export const searchKb = async ({
           where: [
             ['status', ModelDataStatusEnum.ready],
             'AND',
-            ['model_id', model._id],
+            `kb_id IN (${model.chat.relatedKbs.map((item) => `'${item}'`).join(',')})`,
             'AND',
             `vector <=> '[${promptVector}]' < ${similarity}`
           ],

src/service/utils/auth.ts

@@ -34,6 +34,13 @@ export const authToken = (req: NextApiRequest): Promise<string> => {
   });
 };
 
+export const getOpenAiKey = () => {
+  // 纯字符串类型
+  const keys = process.env.OPENAIKEY?.split(',') || [];
+  const i = Math.floor(Math.random() * keys.length);
+  return keys[i] || (process.env.OPENAIKEY as string);
+};
+
 /* 获取 api 请求的 key */
 export const getApiKey = async ({
   model,
@@ -52,7 +59,7 @@ export const getApiKey = async ({
   const keyMap = {
     [OpenAiChatEnum.GPT35]: {
       userOpenAiKey: user.openaiKey || '',
-      systemAuthKey: process.env.OPENAIKEY as string
+      systemAuthKey: getOpenAiKey() as string
     },
     [OpenAiChatEnum.GPT4]: {
       userOpenAiKey: user.openaiKey || '',

src/service/utils/chat/openai.ts

@@ -7,6 +7,7 @@ import { adaptChatItem_openAI } from '@/utils/chat/openai';
 import { modelToolMap } from '@/utils/chat';
 import { ChatCompletionType, ChatContextFilter, StreamResponseType } from './index';
 import { ChatRoleEnum } from '@/constants/chat';
+import { getOpenAiKey } from '../auth';
 
 export const getOpenAIApi = (apiKey: string) => {
   const configuration = new Configuration({
@@ -27,7 +28,7 @@ export const openaiCreateEmbedding = async ({
   userId: string;
   textArr: string[];
 }) => {
-  const systemAuthKey = process.env.OPENAIKEY as string;
+  const systemAuthKey = getOpenAiKey();
 
   // 获取 chatAPI
   const chatAPI = getOpenAIApi(userOpenAiKey || systemAuthKey);