perf: vector generate (#1748)

projects/app/src/pages/api/core/dataset/training/rebuildEmbedding.ts

@@ -81,14 +81,18 @@ async function handler(
   });
 
   // get 10 init dataset.data
-  const arr = new Array(10).fill(0);
+  const max = global.systemEnv?.vectorMaxProcess || 10;
+  const arr = new Array(max * 2).fill(0);
+
   for await (const _ of arr) {
-    await mongoSessionRun(async (session) => {
+    try {
+      const hasNext = await mongoSessionRun(async (session) => {
+        // get next dataset.data
         const data = await MongoDatasetData.findOneAndUpdate(
           {
+            rebuilding: true,
             teamId,
-          datasetId,
-          rebuilding: true
+            datasetId
           },
           {
             $unset: {
@@ -124,7 +128,14 @@ async function handler(
             }
           );
         }
+
+        return !!data;
       });
+
+      if (!hasNext) {
+        break;
+      }
+    } catch (error) {}
   }
 
   return {};

projects/app/src/service/events/generateVector.ts

@@ -158,43 +158,8 @@ const rebuildData = async ({
 
   const deleteVectorIdList = mongoData.indexes.map((index) => index.dataId);
 
-  const { tokens } = await mongoSessionRun(async (session) => {
-    // update vector, update dataset.data rebuilding status, delete data from training
-    const updateResult = await Promise.all(
-      mongoData.indexes.map(async (index, i) => {
-        const result = await insertDatasetDataVector({
-          query: index.text,
-          model: getVectorModel(trainingData.model),
-          teamId: mongoData.teamId,
-          datasetId: mongoData.datasetId,
-          collectionId: mongoData.collectionId
-        });
-        mongoData.indexes[i].dataId = result.insertId;
-        return result;
-      })
-    );
-
-    // Ensure that the training data is deleted after the Mongo update is successful
-    await mongoData.save({ session });
-    await trainingData.deleteOne({ session });
-
-    // delete old vector
-    await deleteDatasetDataVector({
-      teamId: mongoData.teamId,
-      idList: deleteVectorIdList
-    });
-
-    return {
-      tokens: updateResult.reduce((acc, cur) => acc + cur.tokens, 0)
-    };
-  });
-
-  // find next data insert to training queue
-  const arr = new Array(5).fill(0);
-
-  for await (const _ of arr) {
-    try {
-      const hasNextData = await mongoSessionRun(async (session) => {
+  // Find next rebuilding data to insert training queue
+  await mongoSessionRun(async (session) => {
     // get new mongoData insert to training
     const newRebuildingData = await MongoDatasetData.findOneAndUpdate(
       {
@@ -232,15 +197,39 @@ const rebuildData = async ({
         { session }
       );
     }
-
-        return !!newRebuildingData;
   });
 
-      if (!hasNextData) {
-        break;
-      }
-    } catch (error) {}
-  }
+  // update vector, update dataset_data rebuilding status, delete data from training
+  // 1. Insert new vector to dataset_data
+  const updateResult = await Promise.all(
+    mongoData.indexes.map(async (index, i) => {
+      const result = await insertDatasetDataVector({
+        query: index.text,
+        model: getVectorModel(trainingData.model),
+        teamId: mongoData.teamId,
+        datasetId: mongoData.datasetId,
+        collectionId: mongoData.collectionId
+      });
+      mongoData.indexes[i].dataId = result.insertId;
+      return result;
+    })
+  );
+  const { tokens } = await mongoSessionRun(async (session) => {
+    // 2. Ensure that the training data is deleted after the Mongo update is successful
+    await mongoData.save({ session });
+    // 3. Delete the training data
+    await trainingData.deleteOne({ session });
+
+    // 4. Delete old vector
+    await deleteDatasetDataVector({
+      teamId: mongoData.teamId,
+      idList: deleteVectorIdList
+    });
+
+    return {
+      tokens: updateResult.reduce((acc, cur) => acc + cur.tokens, 0)
+    };
+  });
 
   return { tokens };
 };