From a9db5b57c54be613684a07bc8e148f35874c1d12 Mon Sep 17 00:00:00 2001
From: Archer <545436317@qq.com>
Date: Mon, 18 Nov 2024 20:57:03 +0800
Subject: [PATCH] feat: add training retry time (#3187)

* feat: add training retry time

* remoce log
---
 .../service/core/dataset/training/schema.ts   |  7 +++-
 .../app/src/pages/api/core/dataset/update.ts  | 35 ++++++++++++++++++-
 projects/app/src/service/events/generateQA.ts |  8 +++--
 .../app/src/service/events/generateVector.ts  |  6 ++--
 4 files changed, 49 insertions(+), 7 deletions(-)

diff --git a/packages/service/core/dataset/training/schema.ts b/packages/service/core/dataset/training/schema.ts
index ce8724878a2a..48e01613a7e8 100644
--- a/packages/service/core/dataset/training/schema.ts
+++ b/packages/service/core/dataset/training/schema.ts
@@ -51,6 +51,11 @@ const TrainingDataSchema = new Schema({
     type: Date,
     default: () => new Date('2000/1/1')
   },
+  retryCount: {
+    type: Number,
+    default: 5
+  },
+
   model: {
     // ai model
     type: String,
@@ -97,7 +102,7 @@ try {
   // lock training data(teamId); delete training data
   TrainingDataSchema.index({ teamId: 1, datasetId: 1 });
   // get training data and sort
-  TrainingDataSchema.index({ mode: 1, lockTime: 1, weight: -1 });
+  TrainingDataSchema.index({ mode: 1, retryCount: 1, lockTime: 1, weight: -1 });
   TrainingDataSchema.index({ expireAt: 1 }, { expireAfterSeconds: 7 * 24 * 60 * 60 }); // 7 days
 } catch (error) {
   console.log(error);
diff --git a/projects/app/src/pages/api/core/dataset/update.ts b/projects/app/src/pages/api/core/dataset/update.ts
index 26b0c6dcb8da..b018e8bea98f 100644
--- a/projects/app/src/pages/api/core/dataset/update.ts
+++ b/projects/app/src/pages/api/core/dataset/update.ts
@@ -9,7 +9,7 @@ import {
 } from '@fastgpt/global/support/permission/constant';
 import { CommonErrEnum } from '@fastgpt/global/common/error/code/common';
 import type { ApiRequestProps, ApiResponseType } from '@fastgpt/service/type/next';
-import { DatasetTypeEnum } from '@fastgpt/global/core/dataset/constants';
+import { DatasetTypeEnum, TrainingModeEnum } from '@fastgpt/global/core/dataset/constants';
 import { ClientSession } from 'mongoose';
 import { parseParentIdInMongo } from '@fastgpt/global/common/parentFolder/utils';
 import { mongoSessionRun } from '@fastgpt/service/common/mongo/sessionRun';
@@ -21,6 +21,7 @@ import {
 import { authUserPer } from '@fastgpt/service/support/permission/user/auth';
 import { TeamWritePermissionVal } from '@fastgpt/global/support/permission/user/constant';
 import { DatasetErrEnum } from '@fastgpt/global/common/error/code/dataset';
+import { MongoDatasetTraining } from '@fastgpt/service/core/dataset/training/schema';
 
 export type DatasetUpdateQuery = {};
 export type DatasetUpdateResponse = any;
@@ -84,6 +85,12 @@ async function handler(
 
   const isFolder = dataset.type === DatasetTypeEnum.folder;
 
+  updateTraining({
+    teamId: dataset.teamId,
+    datasetId: id,
+    agentModel: agentModel?.model
+  });
+
   const onUpdate = async (session?: ClientSession) => {
     await MongoDataset.findByIdAndUpdate(
       id,
@@ -137,3 +144,29 @@ async function handler(
   }
 }
 export default NextAPI(handler);
+
+async function updateTraining({
+  teamId,
+  datasetId,
+  agentModel
+}: {
+  teamId: string;
+  datasetId: string;
+  agentModel?: string;
+}) {
+  if (!agentModel) return;
+
+  await MongoDatasetTraining.updateMany(
+    {
+      teamId,
+      datasetId,
+      mode: { $in: [TrainingModeEnum.qa, TrainingModeEnum.auto] }
+    },
+    {
+      $set: {
+        model: agentModel,
+        retryCount: 5
+      }
+    }
+  );
+}
diff --git a/projects/app/src/service/events/generateQA.ts b/projects/app/src/service/events/generateQA.ts
index aaa53b88cc67..019fb43536e9 100644
--- a/projects/app/src/service/events/generateQA.ts
+++ b/projects/app/src/service/events/generateQA.ts
@@ -39,11 +39,13 @@ export async function generateQA(): Promise<any> {
     try {
       const data = await MongoDatasetTraining.findOneAndUpdate(
         {
-          lockTime: { $lte: addMinutes(new Date(), -6) },
-          mode: TrainingModeEnum.qa
+          mode: TrainingModeEnum.qa,
+          retryCount: { $gte: 0 },
+          lockTime: { $lte: addMinutes(new Date(), -6) }
         },
         {
-          lockTime: new Date()
+          lockTime: new Date(),
+          $inc: { retryCount: -1 }
         }
       )
         .select({
diff --git a/projects/app/src/service/events/generateVector.ts b/projects/app/src/service/events/generateVector.ts
index e550219a15a2..99fe3383b5ac 100644
--- a/projects/app/src/service/events/generateVector.ts
+++ b/projects/app/src/service/events/generateVector.ts
@@ -39,10 +39,12 @@ export async function generateVector(): Promise<any> {
       const data = await MongoDatasetTraining.findOneAndUpdate(
         {
           mode: TrainingModeEnum.chunk,
-          lockTime: { $lte: addMinutes(new Date(), -1) }
+          retryCount: { $gte: 0 },
+          lockTime: { $lte: addMinutes(new Date(), -6) }
         },
         {
-          lockTime: new Date()
+          lockTime: new Date(),
+          $inc: { retryCount: -1 }
         }
       ).select({
         _id: 1,