package service import ( "context" "rag/dao" "rag/model/dto" "rag/common/task" "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 } 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 } // 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, }) } 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 }