xinyin025/geekai
1
0
Fork 0
mirror of https://github.com/yangjian102621/geekai.git synced 2025-12-05 07:46:06 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

use http pull message to page notify

Browse Source
This commit is contained in:
RockYang
2025-03-04 06:54:30 +08:00
parent 668ff70bc1 0ef51714c9
commit 3958e99e4d
23 changed files with 482 additions and 670 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
api/core/types/task.go
View File

@@ -26,7 +26,6 @@ const (
type MjTask struct {
Id uint `json:"id"` // 任务ID
TaskId string `json:"task_id"` // 中转任务ID
ClientId string `json:"client_id"`
ImgArr []string `json:"img_arr"`
Type TaskType `json:"type"`
UserId int `json:"user_id"`
@@ -44,7 +43,6 @@ type MjTask struct {
type SdTask struct {
Id int `json:"id"` // job 数据库ID
Type TaskType `json:"type"`
ClientId string `json:"client_id"`
UserId int `json:"user_id"`
Params SdTaskParams `json:"params"`
RetryCount int `json:"retry_count"`
@@ -52,7 +50,6 @@ type SdTask struct {
}
type SdTaskParams struct {
ClientId string `json:"client_id"` // 客户端ID
TaskId string `json:"task_id"`
Prompt string `json:"prompt"` // 提示词
NegPrompt string `json:"neg_prompt"` // 反向提示词
@@ -73,7 +70,6 @@ type SdTaskParams struct {
// DallTask DALL-E task
type DallTask struct {
ClientId string `json:"client_id"`
ModelId uint `json:"model_id"`
ModelName string `json:"model_name"`
Id uint `json:"id"`
@@ -88,7 +84,6 @@ type DallTask struct {
}
type SunoTask struct {
ClientId string `json:"client_id"`
Id uint `json:"id"`
Channel string `json:"channel"`
UserId int `json:"user_id"`
@@ -113,7 +108,6 @@ const (
)
type VideoTask struct {
ClientId string `json:"client_id"`
Id uint `json:"id"`
Channel string `json:"channel"`
UserId int `json:"user_id"`

1
api/handler/dalle_handler.go
View File

@@ -70,7 +70,6 @@ func (h *DallJobHandler) Image(c *gin.Context) {
idValue, _ := c.Get(types.LoginUserID)
userId := utils.IntValue(utils.InterfaceToString(idValue), 0)
task := types.DallTask{
ClientId: data.ClientId,
UserId: uint(userId),
ModelId: chatModel.Id,
ModelName: chatModel.Value,

5
api/handler/mj_handler.go
View File

@@ -66,7 +66,6 @@ func (h *MidJourneyHandler) preCheck(c *gin.Context) bool {
func (h *MidJourneyHandler) Image(c *gin.Context) {
var data struct {
TaskType string `json:"task_type"`
ClientId string `json:"client_id"`
Prompt string `json:"prompt"`
NegPrompt string `json:"neg_prompt"`
Rate string `json:"rate"`
@@ -153,7 +152,6 @@ func (h *MidJourneyHandler) Image(c *gin.Context) {
return
}
task := types.MjTask{
ClientId: data.ClientId,
TaskId: taskId,
Type: types.TaskType(data.TaskType),
Prompt: data.Prompt,
@@ -207,7 +205,6 @@ func (h *MidJourneyHandler) Image(c *gin.Context) {
type reqVo struct {
Index int `json:"index"`
ClientId string `json:"client_id"`
ChannelId string `json:"channel_id"`
MessageId string `json:"message_id"`
MessageHash string `json:"message_hash"`
@@ -229,7 +226,6 @@ func (h *MidJourneyHandler) Upscale(c *gin.Context) {
userId := utils.IntValue(utils.InterfaceToString(idValue), 0)
taskId, _ := h.snowflake.Next(true)
task := types.MjTask{
ClientId: data.ClientId,
Type: types.TaskUpscale,
UserId: userId,
ChannelId: data.ChannelId,
@@ -286,7 +282,6 @@ func (h *MidJourneyHandler) Variation(c *gin.Context) {
taskId, _ := h.snowflake.Next(true)
task := types.MjTask{
Type: types.TaskVariation,
ClientId: data.ClientId,
UserId: userId,
Index: data.Index,
ChannelId: data.ChannelId,

3
api/handler/sd_handler.go
View File

@@ -112,8 +112,7 @@ func (h *SdJobHandler) Image(c *gin.Context) {
}
task := types.SdTask{
ClientId: data.ClientId,
Type: types.TaskImage,
Type: types.TaskImage,
Params: types.SdTaskParams{
TaskId: taskId,
Prompt: data.Prompt,

2
api/handler/suno_handler.go
View File

@@ -45,7 +45,6 @@ func NewSunoHandler(app *core.AppServer, db *gorm.DB, service *suno.Service, upl
func (h *SunoHandler) Create(c *gin.Context) {
var data struct {
ClientId string `json:"client_id"`
Prompt string `json:"prompt"`
Instrumental bool `json:"instrumental"`
Lyrics string `json:"lyrics"`
@@ -90,7 +89,6 @@ func (h *SunoHandler) Create(c *gin.Context) {
}
}
task := types.SunoTask{
ClientId: data.ClientId,
UserId: int(h.GetLoginUserId(c)),
Type: data.Type,
Title: data.Title,

8
api/handler/video_handler.go
View File

@@ -46,7 +46,6 @@ func NewVideoHandler(app *core.AppServer, db *gorm.DB, service *video.Service, u
func (h *VideoHandler) LumaCreate(c *gin.Context) {
var data struct {
ClientId string `json:"client_id"`
Prompt string `json:"prompt"`
FirstFrameImg string `json:"first_frame_img,omitempty"`
EndFrameImg string `json:"end_frame_img,omitempty"`
@@ -82,7 +81,6 @@ func (h *VideoHandler) LumaCreate(c *gin.Context) {
EndImgURL: data.EndFrameImg,
}
task := types.VideoTask{
ClientId: data.ClientId,
UserId: userId,
Type: types.VideoLuma,
Prompt: data.Prompt,
@@ -124,7 +122,6 @@ func (h *VideoHandler) KeLingCreate(c *gin.Context) {
var data struct {
Channel string `json:"channel"`
ClientId string `json:"client_id"`
TaskType string `json:"task_type"` // 任务类型: text2video/image2video
Model string `json:"model"` // 模型: default/anime
Prompt string `json:"prompt"` // 视频描述
@@ -173,7 +170,6 @@ func (h *VideoHandler) KeLingCreate(c *gin.Context) {
ImageTail: data.ImageTail,
}
task := types.VideoTask{
ClientId: data.ClientId,
UserId: userId,
Type: types.VideoKeLing,
Prompt: data.Prompt,
@@ -218,14 +214,14 @@ func (h *VideoHandler) List(c *gin.Context) {
page := h.GetInt(c, "page", 1)
pageSize := h.GetInt(c, "page_size", 20)
all := h.GetBool(c, "all")
session := h.DB.Session(&gorm.Session{}).Where("user_id", userId)
session := h.DB.Session(&gorm.Session{})
if t != "" {
session = session.Where("type", t)
}
if all {
session = session.Where("publish", 0).Where("progress", 100)
} else {
session = session.Where("user_id", h.GetLoginUserId(c))
session = session.Where("user_id", userId)
}
// 统计总数
var total int64

5
api/main.go
View File

@@ -163,7 +163,6 @@ func main() {
fx.Provide(dalle.NewService),
fx.Invoke(func(s *dalle.Service) {
s.Run()
s.CheckTaskNotify()
s.DownloadImages()
s.CheckTaskStatus()
}),
@@ -182,7 +181,6 @@ func main() {
fx.Invoke(func(s *mj.Service) {
s.Run()
s.SyncTaskProgress()
s.CheckTaskNotify()
s.DownloadImages()
}),
@@ -191,21 +189,18 @@ func main() {
fx.Invoke(func(s *sd.Service, config *types.AppConfig) {
s.Run()
s.CheckTaskStatus()
s.CheckTaskNotify()
}),
fx.Provide(suno.NewService),
fx.Invoke(func(s *suno.Service) {
s.Run()
s.SyncTaskProgress()
s.CheckTaskNotify()
s.DownloadFiles()
}),
fx.Provide(video.NewService),
fx.Invoke(func(s *video.Service) {
s.Run()
s.SyncTaskProgress()
s.CheckTaskNotify()
s.DownloadFiles()
}),
fx.Provide(service.NewUserService),

56
api/service/dalle/service.go
View File

@@ -34,10 +34,8 @@ type Service struct {
db *gorm.DB
uploadManager *oss.UploaderManager
taskQueue *store.RedisQueue
notifyQueue *store.RedisQueue
userService *service.UserService
wsService *service.WebsocketService
clientIds map[uint]string
}
func NewService(db *gorm.DB, manager *oss.UploaderManager, redisCli *redis.Client, userService *service.UserService, wsService *service.WebsocketService) *Service {
@@ -45,11 +43,9 @@ func NewService(db *gorm.DB, manager *oss.UploaderManager, redisCli *redis.Clien
httpClient: req.C().SetTimeout(time.Minute * 3),
db: db,
taskQueue: store.NewRedisQueue("DallE_Task_Queue", redisCli),
notifyQueue: store.NewRedisQueue("DallE_Notify_Queue", redisCli),
wsService: wsService,
uploadManager: manager,
userService: userService,
clientIds: map[uint]string{},
}
}
@@ -60,7 +56,7 @@ func (s *Service) PushTask(task types.DallTask) {
}
func (s *Service) Run() {
// 将数据库中未提交的人物加载到队列
// 将数据库中未提交的任务加载到队列
var jobs []model.DallJob
s.db.Where("progress", 0).Find(&jobs)
for _, v := range jobs {
@@ -84,16 +80,16 @@ func (s *Service) Run() {
continue
}
logger.Infof("handle a new DALL-E task: %+v", task)
s.clientIds[task.Id] = task.ClientId
_, err = s.Image(task, false)
if err != nil {
logger.Errorf("error with image task: %v", err)
s.db.Model(&model.DallJob{Id: task.Id}).UpdateColumns(map[string]interface{}{
"progress": service.FailTaskProgress,
"err_msg": err.Error(),
})
s.notifyQueue.RPush(service.NotifyMessage{ClientId: task.ClientId, UserId: int(task.UserId), JobId: int(task.Id), Message: service.TaskStatusFailed})
}
go func() {
_, err = s.Image(task, false)
if err != nil {
logger.Errorf("error with image task: %v", err)
s.db.Model(&model.DallJob{Id: task.Id}).UpdateColumns(map[string]interface{}{
"progress": service.FailTaskProgress,
"err_msg": err.Error(),
})
}
}()
}
}()
}
@@ -212,10 +208,9 @@ func (s *Service) Image(task types.DallTask, sync bool) (string, error) {
return "", fmt.Errorf("err with update database: %v", err)
}
s.notifyQueue.RPush(service.NotifyMessage{ClientId: task.ClientId, UserId: int(task.UserId), JobId: int(task.Id), Message: service.TaskStatusFailed})
var content string
if sync {
imgURL, err := s.downloadImage(task.Id, int(task.UserId), res.Data[0].Url)
imgURL, err := s.downloadImage(task.Id, res.Data[0].Url)
if err != nil {
return "", fmt.Errorf("error with download image: %v", err)
}
@@ -225,26 +220,6 @@ func (s *Service) Image(task types.DallTask, sync bool) (string, error) {
return content, nil
}
func (s *Service) CheckTaskNotify() {
go func() {
logger.Info("Running DALL-E task notify checking ...")
for {
var message service.NotifyMessage
err := s.notifyQueue.LPop(&message)
if err != nil {
continue
}
logger.Debugf("notify message: %+v", message)
client := s.wsService.Clients.Get(message.ClientId)
if client == nil {
continue
}
utils.SendChannelMsg(client, types.ChDall, message.Message)
}
}()
}
func (s *Service) CheckTaskStatus() {
go func() {
logger.Info("Running DALL-E task status checking ...")
@@ -254,7 +229,7 @@ func (s *Service) CheckTaskStatus() {
s.db.Where("progress < ?", 100).Find(&jobs)
for _, job := range jobs {
// 超时的任务标记为失败
if time.Now().Sub(job.CreatedAt) > time.Minute*10 {
if time.Since(job.CreatedAt) > time.Minute*10 {
job.Progress = service.FailTaskProgress
job.ErrMsg = "任务超时"
s.db.Updates(&job)
@@ -301,7 +276,7 @@ func (s *Service) DownloadImages() {
}
logger.Infof("try to download image: %s", v.OrgURL)
imgURL, err := s.downloadImage(v.Id, int(v.UserId), v.OrgURL)
imgURL, err := s.downloadImage(v.Id, v.OrgURL)
if err != nil {
logger.Error("error with download image: %s, error: %v", imgURL, err)
continue
@@ -316,7 +291,7 @@ func (s *Service) DownloadImages() {
}()
}
func (s *Service) downloadImage(jobId uint, userId int, orgURL string) (string, error) {
func (s *Service) downloadImage(jobId uint, orgURL string) (string, error) {
// sava image
imgURL, err := s.uploadManager.GetUploadHandler().PutUrlFile(orgURL, false)
if err != nil {
@@ -328,6 +303,5 @@ func (s *Service) downloadImage(jobId uint, userId int, orgURL string) (string,
if res.Error != nil {
return "", err
}
s.notifyQueue.RPush(service.NotifyMessage{ClientId: s.clientIds[jobId], UserId: userId, JobId: int(jobId), Message: service.TaskStatusFinished})
return imgURL, nil
}

55
api/service/mj/service.go
View File

@@ -15,10 +15,11 @@ import (
"geekai/store"
"geekai/store/model"
"geekai/utils"
"github.com/go-redis/redis/v8"
"strings"
"time"
"github.com/go-redis/redis/v8"
"gorm.io/gorm"
)
@@ -26,23 +27,19 @@ import (
type Service struct {
client *Client // MJ Client
taskQueue *store.RedisQueue
notifyQueue *store.RedisQueue
db *gorm.DB
wsService *service.WebsocketService
uploaderManager *oss.UploaderManager
userService *service.UserService
clientIds map[uint]string
}
func NewService(redisCli *redis.Client, db *gorm.DB, client *Client, manager *oss.UploaderManager, wsService *service.WebsocketService, userService *service.UserService) *Service {
return &Service{
db: db,
taskQueue: store.NewRedisQueue("MidJourney_Task_Queue", redisCli),
notifyQueue: store.NewRedisQueue("MidJourney_Notify_Queue", redisCli),
client: client,
wsService: wsService,
uploaderManager: manager,
clientIds: map[uint]string{},
userService: userService,
}
}
@@ -59,7 +56,6 @@ func (s *Service) Run() {
continue
}
task.Id = v.Id
s.clientIds[task.Id] = task.ClientId
s.PushTask(task)
}
@@ -96,7 +92,6 @@ func (s *Service) Run() {
if task.Mode == "" {
task.Mode = "fast"
}
s.clientIds[task.Id] = task.ClientId
var job model.MidJourneyJob
tx := s.db.Where("id = ?", task.Id).First(&job)
@@ -139,7 +134,6 @@ func (s *Service) Run() {
// update the task progress
s.db.Updates(&job)
// 任务失败，通知前端
s.notifyQueue.RPush(service.NotifyMessage{ClientId: task.ClientId, UserId: task.UserId, JobId: int(job.Id), Message: service.TaskStatusFailed})
continue
}
logger.Infof("任务提交成功：%+v", res)
@@ -178,24 +172,6 @@ func GetImageHash(action string) string {
return split[len(split)-1]
}
func (s *Service) CheckTaskNotify() {
go func() {
for {
var message service.NotifyMessage
err := s.notifyQueue.LPop(&message)
if err != nil {
continue
}
logger.Debugf("receive a new mj notify message: %+v", message)
client := s.wsService.Clients.Get(message.ClientId)
if client == nil {
continue
}
utils.SendChannelMsg(client, types.ChMj, message.Message)
}
}()
}
func (s *Service) DownloadImages() {
go func() {
var items []model.MidJourneyJob
@@ -228,12 +204,6 @@ func (s *Service) DownloadImages() {
v.ImgURL = imgURL
s.db.Updates(&v)
s.notifyQueue.RPush(service.NotifyMessage{
ClientId: s.clientIds[v.Id],
UserId: v.UserId,
JobId: int(v.Id),
Message: service.TaskStatusFinished})
}
time.Sleep(time.Second * 5)
@@ -259,7 +229,7 @@ func (s *Service) SyncTaskProgress() {
for _, job := range jobs {
// 10 分钟还没完成的任务标记为失败
if time.Now().Sub(job.CreatedAt) > time.Minute*10 {
if time.Since(job.CreatedAt) > time.Minute*10 {
job.Progress = service.FailTaskProgress
job.ErrMsg = "任务超时"
s.db.Updates(&job)
@@ -279,18 +249,12 @@ func (s *Service) SyncTaskProgress() {
"err_msg": task.FailReason,
})
logger.Errorf("task failed: %v", task.FailReason)
s.notifyQueue.RPush(service.NotifyMessage{
ClientId: s.clientIds[job.Id],
UserId: job.UserId,
JobId: int(job.Id),
Message: service.TaskStatusFailed})
continue
}
if len(task.Buttons) > 0 {
job.Hash = GetImageHash(task.Buttons[0].CustomId)
}
oldProgress := job.Progress
job.Progress = utils.IntValue(strings.Replace(task.Progress, "%", "", 1), 0)
if task.ImageUrl != "" {
job.OrgURL = task.ImageUrl
@@ -300,19 +264,6 @@ func (s *Service) SyncTaskProgress() {
logger.Errorf("error with update database: %v", err)
continue
}
// 通知前端更新任务进度
if oldProgress != job.Progress {
message := service.TaskStatusRunning
if job.Progress == 100 {
message = service.TaskStatusFinished
}
s.notifyQueue.RPush(service.NotifyMessage{
ClientId: s.clientIds[job.Id],
UserId: job.UserId,
JobId: int(job.Id),
Message: message})
}
}
// 找出失败的任务，并恢复其扣减算力

41
api/service/sd/service.go
View File

@@ -16,9 +16,10 @@ import (
"geekai/store"
"geekai/store/model"
"geekai/utils"
"github.com/go-redis/redis/v8"
"time"
"github.com/go-redis/redis/v8"
"github.com/imroc/req/v3"
"gorm.io/gorm"
)
@@ -30,7 +31,6 @@ var logger = logger2.GetLogger()
type Service struct {
httpClient *req.Client
taskQueue *store.RedisQueue
notifyQueue *store.RedisQueue
db *gorm.DB
uploadManager *oss.UploaderManager
wsService *service.WebsocketService
@@ -41,7 +41,6 @@ func NewService(db *gorm.DB, manager *oss.UploaderManager, levelDB *store.LevelD
return &Service{
httpClient: req.C(),
taskQueue: store.NewRedisQueue("StableDiffusion_Task_Queue", redisCli),
notifyQueue: store.NewRedisQueue("StableDiffusion_Queue", redisCli),
db: db,
wsService: wsService,
uploadManager: manager,
@@ -102,8 +101,6 @@ func (s *Service) Run() {
"progress": service.FailTaskProgress,
"err_msg": err.Error(),
})
// 通知前端，任务失败
s.notifyQueue.RPush(service.NotifyMessage{ClientId: task.ClientId, UserId: task.UserId, JobId: task.Id, Message: service.TaskStatusFailed})
continue
}
}
@@ -225,15 +222,12 @@ func (s *Service) Txt2Img(task types.SdTask) error {
// task finished
s.db.Model(&model.SdJob{Id: uint(task.Id)}).UpdateColumn("progress", 100)
s.notifyQueue.RPush(service.NotifyMessage{ClientId: task.ClientId, UserId: task.UserId, JobId: task.Id, Message: service.TaskStatusFinished})
return nil
default:
err, resp := s.checkTaskProgress(apiKey)
resp, err := s.checkTaskProgress(apiKey)
// 更新任务进度
if err == nil && resp.Progress > 0 {
s.db.Model(&model.SdJob{Id: uint(task.Id)}).UpdateColumn("progress", int(resp.Progress*100))
// 发送更新状态信号
s.notifyQueue.RPush(service.NotifyMessage{ClientId: task.ClientId, UserId: task.UserId, JobId: task.Id, Message: service.TaskStatusRunning})
}
time.Sleep(time.Second)
}
@@ -242,7 +236,7 @@ func (s *Service) Txt2Img(task types.SdTask) error {
}
// 执行任务
func (s *Service) checkTaskProgress(apiKey model.ApiKey) (error, *TaskProgressResp) {
func (s *Service) checkTaskProgress(apiKey model.ApiKey) (*TaskProgressResp, error) {
apiURL := fmt.Sprintf("%s/sdapi/v1/progress?skip_current_image=false", apiKey.ApiURL)
var res TaskProgressResp
response, err := s.httpClient.R().
@@ -250,13 +244,13 @@ func (s *Service) checkTaskProgress(apiKey model.ApiKey) (error, *TaskProgressRe
SetSuccessResult(&res).
Get(apiURL)
if err != nil {
return err, nil
return nil, err
}
if response.IsErrorState() {
return fmt.Errorf("error http code status: %v", response.Status), nil
return nil, fmt.Errorf("error http code status: %v", response.Status)
}
return nil, &res
return &res, nil
}
func (s *Service) PushTask(task types.SdTask) {
@@ -264,25 +258,6 @@ func (s *Service) PushTask(task types.SdTask) {
s.taskQueue.RPush(task)
}
func (s *Service) CheckTaskNotify() {
go func() {
logger.Info("Running Stable-Diffusion task notify checking ...")
for {
var message service.NotifyMessage
err := s.notifyQueue.LPop(&message)
if err != nil {
continue
}
logger.Debugf("notify message: %+v", message)
client := s.wsService.Clients.Get(message.ClientId)
if client == nil {
continue
}
utils.SendChannelMsg(client, types.ChSd, message.Message)
}
}()
}
// CheckTaskStatus 检查任务状态，自动删除过期或者失败的任务
func (s *Service) CheckTaskStatus() {
go func() {
@@ -297,7 +272,7 @@ func (s *Service) CheckTaskStatus() {
for _, job := range jobs {
// 5 分钟还没完成的任务标记为失败
if time.Now().Sub(job.CreatedAt) > time.Minute*5 {
if time.Since(job.CreatedAt) > time.Minute*5 {
job.Progress = service.FailTaskProgress
job.ErrMsg = "任务超时"
s.db.Updates(&job)

32
api/service/suno/service.go
View File

@@ -18,10 +18,11 @@ import (
"geekai/store"
"geekai/store/model"
"geekai/utils"
"github.com/go-redis/redis/v8"
"io"
"time"
"github.com/go-redis/redis/v8"
"github.com/imroc/req/v3"
"gorm.io/gorm"
)
@@ -35,7 +36,6 @@ type Service struct {
taskQueue *store.RedisQueue
notifyQueue *store.RedisQueue
wsService *service.WebsocketService
clientIds map[string]string
userService *service.UserService
}
@@ -47,7 +47,6 @@ func NewService(db *gorm.DB, manager *oss.UploaderManager, redisCli *redis.Clien
notifyQueue: store.NewRedisQueue("Suno_Notify_Queue", redisCli),
uploadManager: manager,
wsService: wsService,
clientIds: map[string]string{},
userService: userService,
}
}
@@ -70,7 +69,6 @@ func (s *Service) Run() {
}
task.Id = v.Id
s.PushTask(task)
s.clientIds[v.TaskId] = task.ClientId
}
logger.Info("Starting Suno job consumer...")
go func() {
@@ -95,7 +93,6 @@ func (s *Service) Run() {
"err_msg": err.Error(),
"progress": service.FailTaskProgress,
})
s.notifyQueue.RPush(service.NotifyMessage{ClientId: task.ClientId, UserId: task.UserId, JobId: int(task.Id), Message: service.TaskStatusFailed})
continue
}
@@ -104,7 +101,6 @@ func (s *Service) Run() {
"task_id": r.Data,
"channel": r.Channel,
})
s.clientIds[r.Data] = task.ClientId
}
}()
}
@@ -262,27 +258,6 @@ func (s *Service) Upload(task types.SunoTask) (RespVo, error) {
return res, nil
}
func (s *Service) CheckTaskNotify() {
go func() {
logger.Info("Running Suno task notify checking ...")
for {
var message service.NotifyMessage
err := s.notifyQueue.LPop(&message)
if err != nil {
continue
}
logger.Debugf("notify message: %+v", message)
logger.Debugf("client id: %+v", s.wsService.Clients)
client := s.wsService.Clients.Get(message.ClientId)
logger.Debugf("%+v", client)
if client == nil {
continue
}
utils.SendChannelMsg(client, types.ChSuno, message.Message)
}
}()
}
func (s *Service) DownloadFiles() {
go func() {
var items []model.SunoJob
@@ -311,7 +286,6 @@ func (s *Service) DownloadFiles() {
v.AudioURL = audioURL
v.Progress = 100
s.db.Updates(&v)
s.notifyQueue.RPush(service.NotifyMessage{ClientId: s.clientIds[v.TaskId], UserId: v.UserId, JobId: int(v.Id), Message: service.TaskStatusFinished})
}
time.Sleep(time.Second * 10)
@@ -377,12 +351,10 @@ func (s *Service) SyncTaskProgress() {
}
}
tx.Commit()
s.notifyQueue.RPush(service.NotifyMessage{ClientId: s.clientIds[job.TaskId], UserId: job.UserId, JobId: int(job.Id), Message: service.TaskStatusFinished})
} else if task.Data.FailReason != "" {
job.Progress = service.FailTaskProgress
job.ErrMsg = task.Data.FailReason
s.db.Updates(&job)
s.notifyQueue.RPush(service.NotifyMessage{ClientId: s.clientIds[job.TaskId], UserId: job.UserId, JobId: int(job.Id), Message: service.TaskStatusFailed})
}
}

40
api/service/video/video.go
View File

@@ -20,7 +20,6 @@ import (
"geekai/store/model"
"geekai/utils"
"io"
"io/ioutil"
"net/http"
"time"
@@ -37,9 +36,7 @@ type Service struct {
db *gorm.DB
uploadManager *oss.UploaderManager
taskQueue *store.RedisQueue
notifyQueue *store.RedisQueue
wsService *service.WebsocketService
clientIds map[uint]string
userService *service.UserService
}
@@ -48,10 +45,8 @@ func NewService(db *gorm.DB, manager *oss.UploaderManager, redisCli *redis.Clien
httpClient: req.C().SetTimeout(time.Minute * 3),
db: db,
taskQueue: store.NewRedisQueue("Video_Task_Queue", redisCli),
notifyQueue: store.NewRedisQueue("Video_Notify_Queue", redisCli),
wsService: wsService,
uploadManager: manager,
clientIds: map[uint]string{},
userService: userService,
}
}
@@ -74,7 +69,6 @@ func (s *Service) Run() {
}
task.Id = v.Id
s.PushTask(task)
s.clientIds[v.Id] = task.ClientId
}
logger.Info("Starting Video job consumer...")
go func() {
@@ -86,10 +80,6 @@ func (s *Service) Run() {
continue
}
if task.ClientId != "" {
s.clientIds[task.Id] = task.ClientId
}
if task.Type == types.VideoLuma {
// translate prompt
if utils.HasChinese(task.Prompt) {
@@ -112,7 +102,6 @@ func (s *Service) Run() {
if err != nil {
logger.Errorf("update task with error: %v", err)
}
s.notifyQueue.RPush(service.NotifyMessage{ClientId: task.ClientId, UserId: task.UserId, JobId: int(task.Id), Message: service.TaskStatusFailed, Type: types.VideoLuma})
continue
}
@@ -150,7 +139,6 @@ func (s *Service) Run() {
if err != nil {
logger.Errorf("update task with error: %v", err)
}
s.notifyQueue.RPush(service.NotifyMessage{ClientId: task.ClientId, UserId: task.UserId, JobId: int(task.Id), Message: service.TaskStatusFailed, Type: types.VideoKeLing})
continue
}
@@ -170,25 +158,6 @@ func (s *Service) Run() {
}()
}
func (s *Service) CheckTaskNotify() {
go func() {
logger.Info("Running Suno task notify checking ...")
for {
var message service.NotifyMessage
err := s.notifyQueue.LPop(&message)
if err != nil {
continue
}
logger.Debugf("Receive notify message: %+v", message)
client := s.wsService.Clients.Get(message.ClientId)
if client == nil {
continue
}
utils.SendChannelMsg(client, types.ChLuma, message.Message)
}
}()
}
func (s *Service) DownloadFiles() {
go func() {
var items []model.VideoJob
@@ -232,7 +201,6 @@ func (s *Service) DownloadFiles() {
continue
}
s.notifyQueue.RPush(service.NotifyMessage{ClientId: s.clientIds[v.Id], UserId: v.UserId, JobId: int(v.Id), Message: service.TaskStatusFinished, Type: videoTask.Type})
}
time.Sleep(time.Second * 10)
@@ -334,6 +302,12 @@ func (s *Service) SyncTaskProgress() {
logger.Errorf("更新数据库失败：%v", err)
continue
}
} else if task.TaskStatus == "failed" {
// 更新任务信息
s.db.Model(&model.VideoJob{Id: job.Id}).UpdateColumns(map[string]interface{}{
"progress": service.FailTaskProgress,
"err_msg": task.TaskStatusMsg,
})
}
}
@@ -672,7 +646,7 @@ func (s *Service) QueryKeLingTask(taskId string, channel string, action string)
return VideoCallbackData{}, fmt.Errorf("unexpected status code: %d", res.StatusCode)
}
body, err := ioutil.ReadAll(res.Body)
body, err := io.ReadAll(res.Body)
if err != nil {
return VideoCallbackData{}, fmt.Errorf("failed to read response body: %w", err)
}