package service import ( "context" "rag/dao" "rag/model/dto" "rag/common/task" "gitea.com/red-future/common/db/gfdb" "github.com/gogf/gf/v2/database/gdb" "github.com/gogf/gf/v2/frame/g" "github.com/gogf/gf/v2/util/gconv" ) var Task = new(taskService) type taskService struct{} // WriteTaskProgress 写入任务进度(核心方法) func (s *taskService) WriteTaskProgress(ctx context.Context, req *dto.WriteTaskProgressReq) (err error) { t, total, err := dao.Task.Get(ctx, &dto.GetTaskReq{ TaskId: req.TaskId, }) if err != nil { g.Log().Errorf(ctx, "查询任务失败: %v", err) return err } completed := false if total != 0 { taskVO := make([]*dto.TaskVO, 0, total) err = gconv.Struct(t, &taskVO) if err != nil { g.Log().Errorf(ctx, "转换任务失败: %v", err) return err } taskVO = append(taskVO, &dto.TaskVO{ TaskType: req.TaskType, Status: req.Status, }) completed = IsAllSubTasksCompleted(taskVO) } // 1. 查询是否已存在该文档的该类型任务 existTask, _, err := dao.Task.Get(ctx, &dto.GetTaskReq{ TaskId: req.TaskId, TaskType: req.TaskType, }) if err != nil { g.Log().Errorf(ctx, "查询任务失败: %v", err) return err } err = gfdb.DB(ctx).Transaction(ctx, func(ctx context.Context, tx gdb.TX) error { // 2. 如果不存在,则创建新任务 if g.IsEmpty(existTask) { createReq := &dto.CreateTaskReq{ TaskId: req.TaskId, TaskType: req.TaskType, Status: req.Status, Remark: req.Remark, } _, err = dao.Task.Insert(ctx, createReq) } else { // 3. 如果已存在,则更新任务 updateReq := &dto.UpdateTaskReq{ Id: existTask[0].Id, Status: req.Status, Remark: req.Remark, } _, err = dao.Task.Update(ctx, updateReq) if err != nil { g.Log().Errorf(ctx, "更新任务失败: %v", err) return err } } if completed { // 3. 如果已存在,则更新任务 _, err = dao.Task.Update(ctx, &dto.UpdateTaskReq{ TaskId: req.TaskId, Status: task.TaskStatusCompleted, Remark: "文档解析完成", }) } return nil }) return } // IsAllSubTasksCompleted 判断三个子任务是否全部完成 // 参数:传入当前文档的所有子任务列表 func IsAllSubTasksCompleted(subTasks []*dto.TaskVO) bool { // 必须包含 3 种任务类型 hasKeywords := false hasVector := false hasFullText := false for _, t := range subTasks { // 子任务必须是【已完成】状态才计数 if t.Status == task.TaskStatusCompleted { switch t.TaskType { case task.TaskTypeExtractKeywords: hasKeywords = true case task.TaskTypeGenerateVector: hasVector = true case task.TaskTypeFullTextSearch: hasFullText = true } } } // 三个任务全部完成 → 返回true return hasKeywords && hasVector && hasFullText }