feat: file relate kb

44e772f0 · archer · a3c6d680 · 44e772f0 · 44e772f0 · 44e772f0
Commit 44e772f0 authored Sep 04, 2023 by archer
7 changed files
--- a/client/src/pages/api/plugins/file/read.ts
+++ b/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}; charset=${encoding}`);
-    res.setHeader('Content-Type', file.contentType);
    res.setHeader('Cache-Control', 'public, max-age=3600');
    res.setHeader('Content-Disposition', `inline; filename="${encodeURIComponent(file.filename)}"`);

--- a/client/src/pages/api/plugins/file/upload.ts
+++ b/client/src/pages/api/plugins/file/upload.ts
@@ -38,7 +38,7 @@ 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) {
@@ -46,11 +46,22 @@ class UploadModel {
        }
        resolve({
-          // @ts-ignore
+          files:
-          files: req.files?.map((file) => ({
+            // @ts-ignore
-            ...file,
+            req.files?.map((file) => ({
-            originalname: decodeURIComponent(file.originalname)
+              ...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
          }
        })

--- a/client/src/pages/api/plugins/kb/delete.ts
+++ b/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
+    // delete all pg data
-    await App.updateMany(
+    await PgClient.delete(PgTrainingTableName, {
-      {
+      where: [['user_id', userId], 'AND', ['kb_id', id]]
-        userId
+    });
-      },
-      { $pull: { 'chat.relatedKbs': new Types.ObjectId(id) } }
+    // delete related files
-    );
+    const gridFs = new GridFSStorage('dataset', userId);
+    await gridFs.deleteFilesByKbId(id);
    // delete kb data
    await KB.findOneAndDelete({

--- a/client/src/pages/kb/detail/components/Import/FileSelect.tsx
+++ b/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 }) || ''

--- a/client/src/pages/kb/list/index.tsx
+++ b/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();

--- a/client/src/service/lib/gridfs.ts
+++ b/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);

--- a/client/src/utils/file.ts
+++ b/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));
  });