eta_api/models/rag/ai_task.go

package rag

import (
	"database/sql"
	"eta/eta_api/global"
	"eta/eta_api/utils"
	"fmt"
	"time"
)

// AiTask ai这边的任务表
type AiTask struct {
	AiTaskID                int       `gorm:"primaryKey;column:ai_task_id" description:"-"`
	TaskName                string    `gorm:"column:task_name" description:"任务名称"`
	TaskType                string    `gorm:"column:task_type" description:"任务类型"`
	Status                  string    `gorm:"column:status" description:"任务状态"`
	StartTime               time.Time `gorm:"column:start_time" description:"开始时间"`
	EndTime                 time.Time `gorm:"column:end_time" description:"结束时间"`
	CreateTime              time.Time `gorm:"column:create_time" description:"创建时间"`
	UpdateTime              time.Time `gorm:"column:update_time" description:"更新时间"`
	Parameters              string    `gorm:"column:parameters" description:"执行参数"`
	Logs                    string    `gorm:"column:logs" description:"日志"`
	Errormessage            string    `gorm:"column:ErrorMessage" description:"错误信息"`
	Priority                int       `gorm:"column:priority" description:"优先级"`
	RetryCount              int       `gorm:"column:retry_count" description:"重试次数"`
	EstimatedCompletionTime time.Time `gorm:"column:estimated_completion_time" description:"预计完成时间"`
	ActualCompletitonTime   time.Time `gorm:"column:actual_completiton_time" description:"实际完成时间"`
	Remark                  string    `gorm:"column:remark" description:"备注"`
	SysUserID               int       `gorm:"column:sys_user_id" description:"任务创建人id"`
	SysUserRealName         string    `gorm:"column:sys_user_real_name" description:"任务创建人名称"`
}

// TableName get sql table name.获取数据库表名
func (m *AiTask) TableName() string {
	return "ai_task"
}

// AiTaskColumns get sql column name.获取数据库列名
var AiTaskColumns = struct {
	AiTaskID                string
	TaskName                string
	TaskType                string
	Status                  string
	StartTime               string
	EndTime                 string
	CreateTime              string
	UpdateTime              string
	Parameters              string
	Logs                    string
	Errormessage            string
	Priority                string
	RetryCount              string
	EstimatedCompletionTime string
	ActualCompletitonTime   string
	Remark                  string
	SysUserID               string
	SysUserRealName         string
}{
	AiTaskID:                "ai_task_id",
	TaskName:                "task_name",
	TaskType:                "task_type",
	Status:                  "status",
	StartTime:               "start_time",
	EndTime:                 "end_time",
	CreateTime:              "create_time",
	UpdateTime:              "update_time",
	Parameters:              "parameters",
	Logs:                    "logs",
	Errormessage:            "ErrorMessage",
	Priority:                "priority",
	RetryCount:              "retry_count",
	EstimatedCompletionTime: "estimated_completion_time",
	ActualCompletitonTime:   "actual_completiton_time",
	Remark:                  "remark",
	SysUserID:               "sys_user_id",
	SysUserRealName:         "sys_user_real_name",
}

func (m *AiTask) Create() (err error) {
	err = global.DbMap[utils.DbNameAI].Create(&m).Error

	return
}

func (m *AiTask) Update(updateCols []string) (err error) {
	err = global.DbMap[utils.DbNameAI].Select(updateCols).Updates(&m).Error

	return
}

func (m *AiTask) Del() (err error) {
	err = global.DbMap[utils.DbNameAI].Delete(&m).Error

	return
}

func (m *AiTask) GetByID(id int) (item *AiTask, err error) {
	err = global.DbMap[utils.DbNameAI].Where(fmt.Sprintf("%s = ?", AiTaskColumns.AiTaskID), id).First(&item).Error

	return
}

func (m *AiTask) GetByCondition(condition string, pars []interface{}) (item *AiTask, err error) {
	sqlStr := fmt.Sprintf(`SELECT * FROM %s WHERE 1=1 %s`, m.TableName(), condition)
	err = global.DbMap[utils.DbNameAI].Raw(sqlStr, pars...).First(&item).Error

	return
}

func (m *AiTask) GetListByCondition(field, condition string, pars []interface{}, startSize, pageSize int) (items []*AiTask, err error) {
	if field == "" {
		field = "*"
	}
	sqlStr := fmt.Sprintf(`SELECT %s FROM %s WHERE 1=1 %s order by AiTask_id desc LIMIT ?,?`, field, m.TableName(), condition)
	pars = append(pars, startSize, pageSize)
	err = global.DbMap[utils.DbNameAI].Raw(sqlStr, pars...).Find(&items).Error

	return
}

func (m *AiTask) GetCountByCondition(condition string, pars []interface{}) (total int, err error) {
	var intNull sql.NullInt64
	sqlStr := fmt.Sprintf(`SELECT COUNT(1) total FROM %s WHERE 1=1 %s`, m.TableName(), condition)
	err = global.DbMap[utils.DbNameAI].Raw(sqlStr, pars...).Scan(&intNull).Error
	if err == nil && intNull.Valid {
		total = int(intNull.Int64)
	}

	return
}

// AddAiTask
// @Description: 添加Ai模块的任务
// @author: Roc
// @datetime 2025-04-16 16:55:36
// @param aiTask *AiTask
// @param aiRecordList []*AiTaskRecord
// @return err error
func AddAiTask(aiTask *AiTask, aiRecordList []*AiTaskRecord) (err error) {
	to := global.DbMap[utils.DbNameAI].Begin()
	defer func() {
		if err != nil {
			_ = to.Rollback()
		} else {
			_ = to.Commit()
		}
	}()

	err = to.Create(aiTask).Error
	if err != nil {
		return
	}

	for _, aiTaskRecord := range aiRecordList {
		aiTaskRecord.AiTaskID = aiTask.AiTaskID
	}

	err = to.CreateInBatches(aiRecordList, utils.MultiAddNum).Error
	if err != nil {
		return
	}

	return
}