Fix export dataset (#1436)

* fix: export dataset * remove file buffer
labring · May 10, 2024 · d70efe1 · d70efe1
1 parent 435b2fb
commit d70efe1
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 31 deletions.
diff --git a/projects/app/src/pages/api/core/dataset/collection/create/file.ts b/projects/app/src/pages/api/core/dataset/collection/create/file.ts
@@ -1,10 +1,7 @@
 import type { NextApiRequest, NextApiResponse } from 'next';
 import { jsonRes } from '@fastgpt/service/common/response';
 import { connectToDatabase } from '@/service/mongo';
-import {
-  delFileByFileIdList,
-  readFileContentFromMongo
-} from '@fastgpt/service/common/file/gridfs/controller';
+import { readFileContentFromMongo } from '@fastgpt/service/common/file/gridfs/controller';
 import { authDataset } from '@fastgpt/service/support/permission/auth/dataset';
 import { FileIdCreateDatasetCollectionParams } from '@fastgpt/global/core/dataset/api';
 import { createOneCollection } from '@fastgpt/service/core/dataset/collection/controller';
@@ -24,6 +21,7 @@ import { UsageSourceEnum } from '@fastgpt/global/support/wallet/usage/constants'
 import { getLLMModel, getVectorModel } from '@fastgpt/service/core/ai/model';
 import { hashStr } from '@fastgpt/global/common/string/tools';
 import { startTrainingQueue } from '@/service/core/dataset/training/utils';
+import { MongoRawTextBuffer } from '@fastgpt/service/common/buffer/rawText/schema';
 
 export default async function handler(req: NextApiRequest, res: NextApiResponse<any>) {
   const {
@@ -139,6 +137,9 @@ export default async function handler(req: NextApiRequest, res: NextApiResponse<
       return collectionId;
     });
 
+    // remove buffer
+    await MongoRawTextBuffer.deleteOne({ sourceId: fileId });
+
     startTrainingQueue(true);
 
     jsonRes(res);

diff --git a/projects/app/src/pages/api/support/user/team/limit/exportDatasetLimit.ts b/projects/app/src/pages/api/support/user/team/limit/exportDatasetLimit.ts
@@ -1,34 +1,24 @@
 import type { NextApiRequest, NextApiResponse } from 'next';
-import { jsonRes } from '@fastgpt/service/common/response';
-import { connectToDatabase } from '@/service/mongo';
 import { authDataset } from '@fastgpt/service/support/permission/auth/dataset';
 import { checkExportDatasetLimit } from '@fastgpt/service/support/user/utils';
+import { NextAPI } from '@/service/middle/entry';
 
-export default async function handler(req: NextApiRequest, res: NextApiResponse<any>) {
-  try {
-    await connectToDatabase();
-    const { datasetId } = req.query as {
-      datasetId: string;
-    };
+async function handler(req: NextApiRequest, res: NextApiResponse<any>) {
+  const { datasetId } = req.query as {
+    datasetId: string;
+  };
 
-    if (!datasetId) {
-      throw new Error('datasetId is required');
-    }
-
-    // 凭证校验
-    const { teamId } = await authDataset({ req, authToken: true, datasetId, per: 'w' });
+  if (!datasetId) {
+    throw new Error('datasetId is required');
+  }
 
-    await checkExportDatasetLimit({
-      teamId,
-      limitMinutes: global.feConfigs?.limit?.exportDatasetLimitMinutes
-    });
+  // 凭证校验
+  const { teamId } = await authDataset({ req, authToken: true, datasetId, per: 'w' });
 
-    jsonRes(res);
-  } catch (err) {
-    res.status(500);
-    jsonRes(res, {
-      code: 500,
-      error: err
-    });
-  }
+  await checkExportDatasetLimit({
+    teamId,
+    limitMinutes: global.feConfigs?.limit?.exportDatasetLimitMinutes
+  });
 }
+
+export default NextAPI(handler);
diff --git a/projects/app/src/service/middle/entry.ts b/projects/app/src/service/middle/entry.ts
@@ -13,7 +13,8 @@ export const NextAPI = (...args: NextApiHandler[]): NextApiHandler => {
         response = await handler(req, res);
       }
 
-      if (!res.writableFinished) {
+      const contentType = res.getHeader('Content-Type');
+      if ((!contentType || contentType === 'application/json') && !res.writableFinished) {
         return jsonRes(res, {
           code: 200,
           data: response