geekai/api/service/video/video_service.go

package video

// * +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
// * Copyright 2023 The Geek-AI Authors. All rights reserved.
// * Use of this source code is governed by a Apache-2.0 license
// * that can be found in the LICENSE file.
// * @Author yangjian102621@163.com
// * +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++

import (
	"encoding/json"
	"fmt"
	"geekai/core/types"
	"geekai/log"
	"geekai/service"
	"geekai/service/oss"
	"geekai/service/video/adapters"
	"geekai/store"
	"geekai/store/model"
	"geekai/utils"
	"time"

	"github.com/go-redis/redis/v8"
	"gorm.io/gorm"
)

var logger = log.GetLogger()

type Service struct {
	db            *gorm.DB
	uploadManager *oss.UploaderManager
	taskQueue     *store.RedisQueue
	userService   *service.UserService
	adapters      map[string]adapters.VideoAdapter // provider -> adapter
}

func NewService(db *gorm.DB, manager *oss.UploaderManager, redisCli *redis.Client, userService *service.UserService) *Service {
	service := &Service{
		db:            db,
		taskQueue:     store.NewRedisQueue("Video_Task_Queue", redisCli),
		uploadManager: manager,
		userService:   userService,
		adapters:      make(map[string]VideoAdapter),
	}

	// 注册所有适配器
	service.registerAdapters()

	return service
}

// VideoAdapter 类型别名，指向 adapters.VideoAdapter
type VideoAdapter = adapters.VideoAdapter

// registerAdapters 注册所有视频生成适配器
func (s *Service) registerAdapters() {
	// 注册 Veo 适配器
	veoAdapter := adapters.NewVeoAdapter(s.db)
	s.adapters[veoAdapter.GetProvider()] = veoAdapter

	// 注册 Sora 适配器
	soraAdapter := adapters.NewSoraAdapter(s.db)
	s.adapters[soraAdapter.GetProvider()] = soraAdapter

	// 注册 Luma 适配器
	lumaAdapter := adapters.NewLumaAdapter(s.db)
	s.adapters[lumaAdapter.GetProvider()] = lumaAdapter

	// 注册可灵适配器
	kelingAdapter := adapters.NewKelingAdapter(s.db)
	s.adapters[kelingAdapter.GetProvider()] = kelingAdapter

	// 注册 MiniMax 适配器
	minimaxAdapter := adapters.NewMiniMaxAdapter(s.db)
	s.adapters[minimaxAdapter.GetProvider()] = minimaxAdapter

	// 注册 Wan 适配器
	wanAdapter := adapters.NewWanAdapter(s.db)
	s.adapters[wanAdapter.GetProvider()] = wanAdapter

	// 注册 Doubao 适配器
	doubaoAdapter := adapters.NewDoubaoAdapter(s.db)
	s.adapters[doubaoAdapter.GetProvider()] = doubaoAdapter
}

// getAdapter 获取指定 provider 的适配器
func (s *Service) getAdapter(provider string) (adapters.VideoAdapter, error) {
	adapter, ok := s.adapters[provider]
	if !ok {
		return nil, fmt.Errorf("不支持的视频生成服务提供商: %s", provider)
	}
	return adapter, nil
}

// getVideoConfig 获取视频配置
func (s *Service) getVideoConfig() (*types.VideoConfig, error) {
	return GetVideoConfig(s.db)
}

// CreateTask 统一的创建任务方法
func (s *Service) CreateTask(task types.VideoTask) (adapters.CreateTaskResponse, error) {
	// 获取适配器
	adapter, err := s.getAdapter(task.Type)
	if err != nil {
		return adapters.CreateTaskResponse{}, err
	}

	// 获取视频配置
	videoConfig, err := s.getVideoConfig()
	if err != nil {
		return adapters.CreateTaskResponse{}, err
	}

	// 调用适配器创建任务
	return adapter.CreateTask(task, videoConfig)
}

// QueryTask 统一的查询任务方法
func (s *Service) QueryTask(provider string, taskId string, channel string, modelKey string) (adapters.QueryTaskResponse, error) {
	// 获取适配器
	adapter, err := s.getAdapter(provider)
	if err != nil {
		return adapters.QueryTaskResponse{}, err
	}

	// 获取视频配置
	videoConfig, err := s.getVideoConfig()
	if err != nil {
		return adapters.QueryTaskResponse{}, err
	}

	// 调用适配器查询任务
	return adapter.QueryTask(taskId, channel, videoConfig)
}

func (s *Service) PushTask(task types.VideoTask) {
	logger.Infof("[video] push task to queue jobId=%d type=%s", task.Id, task.Type)
	if err := s.taskQueue.RPush(task); err != nil {
		logger.Errorf("[video] push task to queue failed jobId=%d: %v", task.Id, err)
	}
}

func (s *Service) Run() {
	// 将数据库中未提交的任务加载到队列
	var jobs []model.VideoJob
	s.db.Where("task_id", "").Where("progress", 0).Find(&jobs)
	for _, v := range jobs {
		var task types.VideoTask
		err := utils.JsonDecode(v.Params, &task)
		if err != nil {
			logger.Errorf("decode task info with error: %v", err)
			continue
		}
		task.Id = v.Id
		s.PushTask(task)
	}
	logger.Infof("[video] job consumer started, loaded %d pending jobs from DB", len(jobs))
	go func() {
		for {
			var task types.VideoTask
			err := s.taskQueue.LPop(&task)
			if err != nil {
				logger.Errorf("taking task with error: %v", err)
				continue
			}

			logger.Debugf("[video] submitting task jobId=%d type=%s prompt=%q", task.Id, task.Type, task.Prompt)
			r, err := s.CreateTask(task)
			if err != nil {
				logger.Errorf("[video] submit failed jobId=%d type=%s: %v", task.Id, task.Type, err)
				err = s.db.Model(&model.VideoJob{Id: task.Id}).UpdateColumns(map[string]interface{}{
					"err_msg": err.Error(),
					"status":  types.VideoStatusFailed,
				}).Error
				if err != nil {
					logger.Errorf("update task with error: %v", err)
				}
				continue
			}

			logger.Infof("[video] submit success jobId=%d type=%s taskId=%s channel=%s", task.Id, task.Type, r.TaskId, r.Channel)
			err = s.db.Model(&model.VideoJob{Id: task.Id}).UpdateColumns(map[string]interface{}{
				"task_id": r.TaskId,
				"channel": r.Channel,
				"status":  types.VideoStatusPending,
			}).Error
			if err != nil {
				logger.Errorf("update task with error: %v", err)
				s.PushTask(task)
			}

		}
	}()
}

func (s *Service) DownloadFiles() {
	go func() {
		var items []model.VideoJob
		logger.Info("[video] download files started")
		for {
			err := s.db.Where("status", types.VideoStatusDownloading).Find(&items).Error
			if err != nil {
				logger.Errorf("get downloading tasks with error: %v", err)
				continue
			}

			for _, v := range items {
				if v.VideoURL == "" {
					continue
				}

				logger.Infof("try download video: %s", v.VideoURL)
				videoURL, err := s.uploadManager.GetUploadHandler().PutUrlFile(v.VideoURL, ".mp4", true)
				if err != nil {
					logger.Errorf("download video with error: %v", err)
					continue
				}
				logger.Infof("download video success: %s", videoURL)
				s.db.Model(&model.VideoJob{Id: v.Id}).UpdateColumns(map[string]any{
					"video_url": videoURL,
					"status":    types.VideoStatusSuccess,
					"progress":  100,
				})
			}

			time.Sleep(time.Second * 10)
		}
	}()
}

// SyncTaskProgress 异步拉取任务
func (s *Service) SyncTaskProgress() {
	go func() {
		logger.Info("[video] task status poller started")
		var jobs []model.VideoJob
		for {
			res := s.db.Where("status IN ?", []string{types.VideoStatusInProgress, types.VideoStatusPending}).Where("task_id <> ?", "").Find(&jobs)
			if res.Error != nil {
				continue
			}
			if len(jobs) > 0 {
				logger.Infof("[video] polling task status, in_progress count=%d", len(jobs))
			}

			for _, job := range jobs {
				// 检查任务是否超时（超过 2 小时）
				if time.Since(job.CreatedAt) > 2*time.Hour {
					logger.Warnf("[video] task timeout jobId=%d taskId=%s created_at=%s", job.Id, job.TaskId, job.CreatedAt.Format(time.RFC3339))
					err := s.db.Model(&model.VideoJob{Id: job.Id}).UpdateColumns(map[string]any{
						"status":  types.VideoStatusFailed,
						"err_msg": "任务超时",
					}).Error
					if err != nil {
						logger.Errorf("[video] update timeout task failed jobId=%d: %v", job.Id, err)
					}
					continue
				}

				modelKey := ""
				var videoTask types.VideoTask
				if err := json.Unmarshal([]byte(job.Params), &videoTask); err == nil {
					if paramsMap, ok := videoTask.Params.(map[string]any); ok {
						if model, ok := paramsMap["model"].(string); ok {
							modelKey = model
						}
					}
				}

				logger.Debugf("[video] querying task jobId=%d taskId=%s provider=%s", job.Id, job.TaskId, job.Type)
				task, err := s.QueryTask(job.Type, job.TaskId, job.Channel, modelKey)
				if err != nil {
					logger.Errorf("[video] query failed jobId=%d taskId=%s: %v", job.Id, job.TaskId, err)
					// 更新任务信息
					s.db.Model(&model.VideoJob{Id: job.Id}).UpdateColumns(map[string]any{
						"status":  types.VideoStatusFailed,
						"err_msg": err.Error(),
					})
					continue
				}

				logger.Debugf("[video] task status jobId=%d taskId=%s status=%s", job.Id, job.TaskId, task.Status)
				logger.Debugf("[video] output=%s", task.Output)

				if task.Status == types.VideoStatusSuccess {
					data := map[string]any{
						"status":   types.VideoStatusDownloading,
						"progress": 100,
						"output":   task.Output,
					}
					if task.VideoURL != "" {
						data["video_url"] = task.VideoURL
					}
					err = s.db.Model(&model.VideoJob{Id: job.Id}).UpdateColumns(data).Error
					if err != nil {
						logger.Errorf("更新数据库失败：%v", err)
						continue
					}
					logger.Infof("[video] task completed jobId=%d taskId=%s", job.Id, job.TaskId)
				} else if task.Status == "failed" {
					logger.Warnf("[video] task failed jobId=%d taskId=%s err=%s", job.Id, job.TaskId, task.ErrMsg)
					s.db.Model(&model.VideoJob{Id: job.Id}).UpdateColumns(map[string]any{
						"status":  types.VideoStatusFailed,
						"err_msg": task.ErrMsg,
					})
				} else {
					s.db.Model(&model.VideoJob{Id: job.Id}).UpdateColumns(map[string]any{
						"status":   task.Status,
						"progress": task.Progress,
					})
					continue
				}

			}

			// 找出失败的任务，并恢复其扣减算力
			s.db.Select("id", "user_id", "power", "task_id", "err_msg", "type").
				Where("status", types.VideoStatusFailed).Where("power > ?", 0).Find(&jobs)
			for _, job := range jobs {
				err := s.userService.IncreasePower(job.UserId, job.Power, model.PowerLog{
					Type:   types.PowerRefund,
					Model:  job.Type,
					Remark: fmt.Sprintf("%s 任务失败，退回算力。任务ID：%s，Err:%s", job.Type, job.TaskId, job.ErrMsg),
				})
				if err != nil {
					continue
				}
				// 更新任务状态
				s.db.Model(&job).UpdateColumn("power", 0)
			}
			time.Sleep(time.Second * 10)
		}
	}()
}