116 lines
2.8 KiB
Go
116 lines
2.8 KiB
Go
package service
|
|
|
|
import (
|
|
"context"
|
|
"rag/consts/task"
|
|
"rag/dao"
|
|
"rag/model/dto"
|
|
|
|
"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
|
|
}
|