doc_ai_backed/internal/service/task.go

package service

import (
	"context"
	"errors"
	"fmt"
	"strings"

	"gitea.com/bitwsd/document_ai/internal/model/task"
	"gitea.com/bitwsd/document_ai/internal/storage/dao"
	"gitea.com/bitwsd/document_ai/pkg/log"
	"gitea.com/bitwsd/document_ai/pkg/oss"
)

type TaskService struct {
	recognitionTaskDao   *dao.RecognitionTaskDao
	evaluateTaskDao      *dao.EvaluateTaskDao
	recognitionResultDao *dao.RecognitionResultDao
}

func NewTaskService() *TaskService {
	return &TaskService{
		recognitionTaskDao:   dao.NewRecognitionTaskDao(),
		evaluateTaskDao:      dao.NewEvaluateTaskDao(),
		recognitionResultDao: dao.NewRecognitionResultDao(),
	}
}

func (svc *TaskService) EvaluateTask(ctx context.Context, req *task.EvaluateTaskRequest) error {
	task, err := svc.recognitionTaskDao.GetByTaskNo(dao.DB.WithContext(ctx), req.TaskNo)
	if err != nil {
		log.Error(ctx, "func", "EvaluateTask", "msg", "get task by task no failed", "error", err)
		return err
	}

	if task == nil {
		log.Error(ctx, "func", "EvaluateTask", "msg", "task not found")
		return errors.New("task not found")
	}
	if task.Status != dao.TaskStatusCompleted {
		log.Error(ctx, "func", "EvaluateTask", "msg", "task not finished")
		return errors.New("task not finished")
	}

	evaluateTask := &dao.EvaluateTask{
		TaskID:    task.ID,
		Satisfied: req.Satisfied,
		Feedback:  req.Feedback,
		Comment:   strings.Join(req.Suggestion, ","),
	}
	err = svc.evaluateTaskDao.Create(dao.DB.WithContext(ctx), evaluateTask)
	if err != nil {
		log.Error(ctx, "func", "EvaluateTask", "msg", "create evaluate task failed", "error", err)
		return err
	}

	return nil
}

func (svc *TaskService) GetTaskList(ctx context.Context, req *task.TaskListRequest) (*task.TaskListResponse, error) {
	tasks, total, err := svc.recognitionTaskDao.GetTaskList(dao.DB.WithContext(ctx), req.UserID, dao.TaskType(req.TaskType), req.Page, req.PageSize)
	if err != nil {
		log.Error(ctx, "func", "GetTaskList", "msg", "get task list failed", "error", err)
		return nil, err
	}

	taskIDs := make([]int64, 0, len(tasks))
	for _, item := range tasks {
		taskIDs = append(taskIDs, item.ID)
	}

	recognitionResults, err := svc.recognitionResultDao.GetByTaskIDs(dao.DB.WithContext(ctx), taskIDs)
	if err != nil {
		log.Error(ctx, "func", "GetTaskList", "msg", "get recognition results failed", "error", err)
		return nil, err
	}

	recognitionResultMap := make(map[int64]*dao.RecognitionResult)
	for _, item := range recognitionResults {
		recognitionResultMap[item.TaskID] = item
	}

	resp := &task.TaskListResponse{
		TaskList: make([]*task.TaskListDTO, 0, len(tasks)),
		Total:    total,
	}
	for _, item := range tasks {
		var latex string
		var markdown string
		recognitionResult := recognitionResultMap[item.ID]
		if recognitionResult != nil {
			latex = recognitionResult.NewContentCodec().GetContent().(string)
			markdown = fmt.Sprintf("$$%s$$", latex)
		}
		originURL, err := oss.GetDownloadURL(ctx, item.FileURL)
		if err != nil {
			log.Error(ctx, "func", "GetTaskList", "msg", "get origin url failed", "error", err)
		}
		resp.TaskList = append(resp.TaskList, &task.TaskListDTO{
			Latex:     latex,
			Markdown:  markdown,
			TaskID:    item.TaskUUID,
			FileName:  item.FileName,
			Status:    int(item.Status),
			OriginURL: originURL,
			TaskType:  item.TaskType.String(),
			CreatedAt: item.CreatedAt.Format("2006-01-02 15:04:05"),
		})
	}
	return resp, nil
}