feat: file relate kb

Author: c121914yu

client/src/pages/api/plugins/file/read.ts

@@ -26,8 +26,7 @@ export default async function handler(req: NextApiRequest, res: NextApiResponse<
 
     const encoding = jschardet.detect(buffer)?.encoding;
 
-    res.setHeader('encoding', encoding);
-    res.setHeader('Content-Type', file.contentType);
+    res.setHeader('Content-Type', `${file.contentType}; charset=${encoding}`);
     res.setHeader('Cache-Control', 'public, max-age=3600');
     res.setHeader('Content-Disposition', `inline; filename="${encodeURIComponent(file.filename)}"`);
 

client/src/pages/api/plugins/file/upload.ts

@@ -38,19 +38,30 @@ class UploadModel {
   }).any();
 
   async doUpload(req: NextApiRequest, res: NextApiResponse) {
-    return new Promise<{ files: FileType[] }>((resolve, reject) => {
+    return new Promise<{ files: FileType[]; metadata: Record<string, any> }>((resolve, reject) => {
       // @ts-ignore
       this.uploader(req, res, (error) => {
         if (error) {
           return reject(error);
         }
 
         resolve({
-          // @ts-ignore
-          files: req.files?.map((file) => ({
-            ...file,
-            originalname: decodeURIComponent(file.originalname)
-          }))
+          files:
+            // @ts-ignore
+            req.files?.map((file) => ({
+              ...file,
+              originalname: decodeURIComponent(file.originalname)
+            })) || [],
+          metadata: (() => {
+            if (!req.body?.metadata) return {};
+            try {
+              return JSON.parse(req.body.metadata);
+            } catch (error) {
+              console.log(error);
+
+              return {};
+            }
+          })()
         });
       });
     });
@@ -64,7 +75,7 @@ export default async function handler(req: NextApiRequest, res: NextApiResponse<
     await connectToDatabase();
     const { userId } = await authUser({ req, authToken: true });
 
-    const { files = [] } = await upload.doUpload(req, res);
+    const { files, metadata } = await upload.doUpload(req, res);
 
     const gridFs = new GridFSStorage('dataset', userId);
 
@@ -74,8 +85,8 @@ export default async function handler(req: NextApiRequest, res: NextApiResponse<
           path: file.path,
           filename: file.originalname,
           metadata: {
+            ...metadata,
             contentType: file.mimetype,
-            encoding: file.encoding,
             userId
           }
         })

client/src/pages/api/plugins/kb/delete.ts

@@ -5,6 +5,7 @@ import { authUser } from '@/service/utils/auth';
 import { PgClient } from '@/service/pg';
 import { Types } from 'mongoose';
 import { PgTrainingTableName } from '@/constants/plugin';
+import { GridFSStorage } from '@/service/lib/gridfs';
 
 export default async function handler(req: NextApiRequest, res: NextApiResponse<any>) {
   try {
@@ -21,24 +22,20 @@ export default async function handler(req: NextApiRequest, res: NextApiResponse<
 
     await connectToDatabase();
 
-    // delete all pg data
-    await PgClient.delete(PgTrainingTableName, {
-      where: [['user_id', userId], 'AND', ['kb_id', id]]
-    });
-
     // delete training data
     await TrainingData.deleteMany({
       userId,
       kbId: id
     });
 
-    // delete related app
-    await App.updateMany(
-      {
-        userId
-      },
-      { $pull: { 'chat.relatedKbs': new Types.ObjectId(id) } }
-    );
+    // delete all pg data
+    await PgClient.delete(PgTrainingTableName, {
+      where: [['user_id', userId], 'AND', ['kb_id', id]]
+    });
+
+    // delete related files
+    const gridFs = new GridFSStorage('dataset', userId);
+    await gridFs.deleteFilesByKbId(id);
 
     // delete kb data
     await KB.findOneAndDelete({

client/src/pages/kb/detail/components/Import/FileSelect.tsx

@@ -19,6 +19,7 @@ import dynamic from 'next/dynamic';
 import MyTooltip from '@/components/MyTooltip';
 import { FetchResultItem, DatasetItemType } from '@/types/plugin';
 import { getErrText } from '@/utils/tools';
+import { useUserStore } from '@/store/user';
 
 const UrlFetchModal = dynamic(() => import('./UrlFetchModal'));
 const CreateFileModal = dynamic(() => import('./CreateFileModal'));
@@ -54,6 +55,7 @@ const FileSelect = ({
   showCreateFile = true,
   ...props
 }: Props) => {
+  const { kbDetail } = useUserStore();
   const { Loading: FileSelectLoading } = useLoading();
   const { t } = useTranslation();
 
@@ -109,7 +111,7 @@ const FileSelect = ({
               }
               return '';
             })(),
-            uploadFiles([file], (percent) => {
+            uploadFiles([file], { kbId: kbDetail._id }, (percent) => {
               if (percent < 100) {
                 setSelectingText(
                   t('file.Uploading', { name: file.name.slice(0, 20), percent }) || ''

client/src/pages/kb/list/index.tsx

@@ -31,7 +31,7 @@ const Kb = () => {
   const { toast } = useToast();
   const { openConfirm, ConfirmModal } = useConfirm({
     title: '删除提示',
-    content: '确认删除该知识库？'
+    content: '确认删除该知识库？知识库相关的文件、记录将永久删除，无法恢复！'
   });
   const { myKbList, loadKbList, setKbList } = useUserStore();
 

client/src/service/lib/gridfs.ts

@@ -79,7 +79,7 @@ export class GridFSStorage {
       return Promise.reject(`file not found`);
     }
 
-    if (String(file.metadata?.userId) !== this.uid) {
+    if (file.metadata?.userId !== this.uid) {
       return Promise.reject(ERROR_ENUM.unAuthFile);
     }
 
@@ -100,6 +100,16 @@ export class GridFSStorage {
     return true;
   }
 
+  async deleteFilesByKbId(kbId: string) {
+    if (!kbId) return;
+    const bucket = this.GridFSBucket();
+    const files = await bucket
+      .find({ ['metadata.kbId']: kbId, ['metadata.userId']: this.uid }, { projection: { _id: 1 } })
+      .toArray();
+
+    return Promise.all(files.map((file) => this.delete(String(file._id))));
+  }
+
   async download(id: string) {
     await this.findAndAuthFile(id);
 

client/src/utils/file.ts

@@ -7,8 +7,13 @@ import { uploadImg, postUploadFiles } from '@/api/system';
 /**
  * upload file to mongo gridfs
  */
-export const uploadFiles = (files: File[], percentListen?: (percent: number) => void) => {
+export const uploadFiles = (
+  files: File[],
+  metadata: Record<string, any> = {},
+  percentListen?: (percent: number) => void
+) => {
   const form = new FormData();
+  form.append('metadata', JSON.stringify(metadata));
   files.forEach((file) => {
     form.append('file', file, encodeURIComponent(file.name));
   });