feat: kb crud

src/service/errorCode.ts

@@ -37,7 +37,8 @@ export const proxyError: Record<string, boolean> = {
 export enum ERROR_ENUM {
   unAuthorization = 'unAuthorization',
   insufficientQuota = 'insufficientQuota',
-  unAuthModel = 'unAuthModel'
+  unAuthModel = 'unAuthModel',
+  unAuthKb = 'unAuthKb'
 }
 export const ERROR_RESPONSE: Record<
   any,
@@ -65,5 +66,11 @@ export const ERROR_RESPONSE: Record<
     statusText: ERROR_ENUM.unAuthModel,
     message: '无权使用该模型',
     data: null
+  },
+  [ERROR_ENUM.unAuthKb]: {
+    code: 512,
+    statusText: ERROR_ENUM.unAuthKb,
+    message: '无权使用该知识库',
+    data: null
   }
 };

src/service/events/generateQA.ts

@@ -5,7 +5,7 @@ import { pushSplitDataBill } from '@/service/events/pushBill';
 import { generateVector } from './generateVector';
 import { openaiError2 } from '../errorCode';
 import { PgClient } from '@/service/pg';
-import { ModelSplitDataSchema } from '@/types/mongoSchema';
+import { SplitDataSchema } from '@/types/mongoSchema';
 import { modelServiceToolMap } from '../utils/chat';
 import { ChatRoleEnum } from '@/constants/chat';
 import { getErrMessage } from '../utils/tools';
@@ -32,7 +32,7 @@ export async function generateQA(next = false): Promise<any> {
       { $sample: { size: 1 } }
     ]);
 
-    const dataItem: ModelSplitDataSchema = data[0];
+    const dataItem: SplitDataSchema = data[0];
 
     if (!dataItem) {
       console.log('没有需要生成 QA 的数据');
@@ -127,14 +127,15 @@ A2:
     const resultList = successResponse.map((item) => item.result).flat();
 
     await Promise.allSettled([
+      // 删掉后5个数据
       SplitData.findByIdAndUpdate(dataItem._id, {
         textList: dataItem.textList.slice(0, -5)
-      }), // 删掉后5个数据
+      }),
       // 生成的内容插入 pg
       PgClient.insert('modelData', {
         values: resultList.map((item) => [
           { key: 'user_id', value: dataItem.userId },
-          { key: 'model_id', value: dataItem.modelId },
+          { key: 'kb_id', value: dataItem.kbId },
           { key: 'q', value: item.q },
           { key: 'a', value: item.a },
           { key: 'status', value: 'waiting' }

src/service/models/kb.ts

@@ -0,0 +1,28 @@
+import { Schema, model, models, Model } from 'mongoose';
+import { kbSchema as SchemaType } from '@/types/mongoSchema';
+
+const kbSchema = new Schema({
+  userId: {
+    type: Schema.Types.ObjectId,
+    ref: 'user',
+    required: true
+  },
+  updateTime: {
+    type: Date,
+    default: () => new Date()
+  },
+  avatar: {
+    type: String,
+    default: '/icon/logo.png'
+  },
+  name: {
+    type: String,
+    required: true
+  },
+  tags: {
+    type: [String],
+    default: []
+  }
+});
+
+export const KB: Model<SchemaType> = models['kb'] || model('kb', kbSchema);

src/service/models/splitData.ts

@@ -1,6 +1,6 @@
 /* 模型的知识库 */
 import { Schema, model, models, Model as MongoModel } from 'mongoose';
-import { ModelSplitDataSchema as SplitDataType } from '@/types/mongoSchema';
+import { SplitDataSchema as SplitDataType } from '@/types/mongoSchema';
 
 const SplitDataSchema = new Schema({
   userId: {
@@ -13,9 +13,9 @@ const SplitDataSchema = new Schema({
     type: String,
     required: true
   },
-  modelId: {
+  kbId: {
     type: Schema.Types.ObjectId,
-    ref: 'model',
+    ref: 'kb',
     required: true
   },
   textList: {

src/service/mongo.ts

@@ -52,3 +52,4 @@ export * from './models/openapi';
 export * from './models/promotionRecord';
 export * from './models/collection';
 export * from './models/shareChat';
+export * from './models/kb';

src/service/utils/auth.ts

@@ -1,7 +1,7 @@
 import type { NextApiRequest } from 'next';
 import jwt from 'jsonwebtoken';
 import cookie from 'cookie';
-import { Chat, Model, OpenApi, User, ShareChat } from '../mongo';
+import { Chat, Model, OpenApi, User, ShareChat, KB } from '../mongo';
 import type { ModelSchema } from '@/types/mongoSchema';
 import type { ChatItemSimpleType } from '@/types/chat';
 import mongoose from 'mongoose';
@@ -129,6 +129,18 @@ export const authModel = async ({
   return { model, showModelDetail: model.share.isShareDetail || userId === String(model.userId) };
 };
 
+// 知识库操作权限
+export const authKb = async ({ kbId, userId }: { kbId: string; userId: string }) => {
+  const kb = await KB.findOne({
+    _id: kbId,
+    userId
+  });
+  if (kb) {
+    return kb;
+  }
+  return Promise.reject(ERROR_ENUM.unAuthKb);
+};
+
 // 获取对话校验
 export const authChat = async ({
   modelId,