xinyin025/geekai
1
0
Fork 0
mirror of https://github.com/yangjian102621/geekai.git synced 2025-11-08 18:23:45 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

refactor: 重构项目，为所有的 AI 工具都引入算力，采用算力统一结算各个工具的调用次数和权限

Browse Source
This commit is contained in:
RockYang
2024-03-12 15:40:44 +08:00
parent 316636f83c
commit e746aafa2f
45 changed files with 438 additions and 298 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
api/service/mj/plus/service.go
View File

@@ -96,12 +96,6 @@ func (s *Service) Run() {
s.db.Updates(&job)
// 任务失败，通知前端
s.notifyQueue.RPush(task.UserId)
// restore img_call quota
if task.Type.String() != types.TaskUpscale.String() {
s.db.Model(&model.User{}).Where("id = ?", task.UserId).UpdateColumn("img_calls", gorm.Expr("img_calls + ?", 1))
}
// TODO: 任务提交失败，加入队列重试
continue
}
logger.Infof("任务提交成功：%+v", res)

43
api/service/mj/pool.go
View File

@@ -191,28 +191,43 @@ func (p *ServicePool) SyncTaskProgress() {
go func() {
var items []model.MidJourneyJob
for {
res := p.db.Where("progress >= ? AND progress < ?", 0, 100).Find(&items)
res := p.db.Where("progress < ?", 100).Find(&items)
if res.Error != nil {
continue
}
for _, v := range items {
// 30 分钟还没完成的任务直接删除
if time.Now().Sub(v.CreatedAt) > time.Minute*30 {
p.db.Delete(&v)
// 非放大任务，退回绘图次数
if v.Type != types.TaskUpscale.String() {
p.db.Model(&model.User{}).Where("id = ?", v.UserId).UpdateColumn("img_calls", gorm.Expr("img_calls + ?", 1))
for _, job := range items {
// 失败或者 30 分钟还没完成的任务删除并退回算力
if time.Now().Sub(job.CreatedAt) > time.Minute*30 || job.Progress == -1 {
p.db.Delete(&job)
// 略过 Upscale 任务
if job.Type != types.TaskUpscale.String() {
continue
}
tx := p.db.Model(&model.User{}).Where("id = ?", job.UserId).UpdateColumn("power", gorm.Expr("power + ?", job.Power))
if tx.Error == nil && tx.RowsAffected > 0 {
var user model.User
p.db.Where("id = ?", job.UserId).First(&user)
p.db.Create(&model.PowerLog{
UserId: user.Id,
Username: user.Username,
Type: types.PowerConsume,
Amount: job.Power,
Balance: user.Power + job.Power,
Mark: types.PowerAdd,
Model: "mid-journey",
Remark: fmt.Sprintf("绘画任务失败，退回算力。任务ID：%s", job.TaskId),
CreatedAt: time.Now(),
})
}
}
if !strings.HasPrefix(job.ChannelId, "mj-service-plus") {
continue
}
if !strings.HasPrefix(v.ChannelId, "mj-service-plus") {
continue
}
if servicePlus := p.getServicePlus(v.ChannelId); servicePlus != nil {
_ = servicePlus.Notify(v)
if servicePlus := p.getServicePlus(job.ChannelId); servicePlus != nil {
_ = servicePlus.Notify(job)
}
}

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

@@ -84,7 +84,7 @@ func (s *Service) Run() {
if err != nil {
logger.Error("绘画任务执行失败：", err.Error())
// update the task progress
s.db.Model(&model.MidJourneyJob{Id: uint(task.Id)}).UpdateColumns(map[string]interface{}{
s.db.Model(&model.MidJourneyJob{Id: task.Id}).UpdateColumns(map[string]interface{}{
"progress": -1,
"err_msg": err.Error(),
})