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

restore use power when removed not finish jobs

Browse Source
This commit is contained in:
RockYang
2024-07-31 16:08:46 +08:00
parent 96f1126d02
commit 1d9d487f0e
25 changed files with 425 additions and 487 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
api/core/app_server.go
View File

@@ -225,7 +225,7 @@ func needLogin(c *gin.Context) bool {
c.Request.URL.Path == "/api/payment/doPay" ||
c.Request.URL.Path == "/api/payment/payWays" ||
c.Request.URL.Path == "/api/suno/client" ||
c.Request.URL.Path == "/api/suno/Detail" ||
c.Request.URL.Path == "/api/suno/detail" ||
c.Request.URL.Path == "/api/suno/play" ||
strings.HasPrefix(c.Request.URL.Path, "/api/test") ||
strings.HasPrefix(c.Request.URL.Path, "/api/user/clogin") ||

1
api/core/types/chat.go
View File

@@ -61,7 +61,6 @@ type ChatSession struct {
type ChatModel struct {
Id uint `json:"id"`
Platform string `json:"platform"`
Name string `json:"name"`
Value string `json:"value"`
Power int `json:"power"`

2
api/handler/chatimpl/chat_handler.go
View File

@@ -468,7 +468,7 @@ func (h *ChatHandler) doRequest(ctx context.Context, req types.ApiRequest, sessi
} else {
client = http.DefaultClient
}
logger.Debugf("Sending %s request, Channel:%s, API KEY:%s, PROXY: %s, Model: %s", session.Model.Platform, apiKey.ApiURL, apiURL, apiKey.ProxyURL, req.Model)
logger.Debugf("Sending %s request, API KEY:%s, PROXY: %s, Model: %s", apiKey.ApiURL, apiURL, apiKey.ProxyURL, req.Model)
request.Header.Set("Authorization", fmt.Sprintf("Bearer %s", apiKey.Value))
// 更新API KEY 最后使用时间
h.DB.Model(&model.ApiKey{}).Where("id", apiKey.Id).UpdateColumn("last_used_at", time.Now().Unix())

52
api/handler/dalle_handler.go
View File

@@ -8,6 +8,7 @@ package handler
// * +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
import (
"fmt"
"geekai/core"
"geekai/core/types"
"geekai/service/dalle"
@@ -16,9 +17,9 @@ import (
"geekai/store/vo"
"geekai/utils"
"geekai/utils/resp"
"net/http"
"github.com/gorilla/websocket"
"net/http"
"time"
"github.com/gin-gonic/gin"
"github.com/go-redis/redis/v8"
@@ -178,7 +179,7 @@ func (h *DallJobHandler) getData(finish bool, userId uint, page int, pageSize in
session := h.DB.Session(&gorm.Session{})
if finish {
session = session.Where("progress = ?", 100).Order("id DESC")
session = session.Where("progress >= ?", 100).Order("id DESC")
} else {
session = session.Where("progress < ?", 100).Order("id ASC")
}
@@ -215,20 +216,51 @@ func (h *DallJobHandler) getData(finish bool, userId uint, page int, pageSize in
// Remove remove task image
func (h *DallJobHandler) Remove(c *gin.Context) {
id := h.GetInt(c, "id", 0)
userId := h.GetInt(c, "user_id", 0)
userId := h.GetLoginUserId(c)
var job model.DallJob
if res := h.DB.Where("id = ? AND user_id = ?", id, userId).First(&job); res.Error != nil {
resp.ERROR(c, "记录不存在")
return
}
// remove job recode
res := h.DB.Delete(&model.DallJob{Id: job.Id})
if res.Error != nil {
resp.ERROR(c, res.Error.Error())
// 删除任务
tx := h.DB.Begin()
if err := tx.Delete(&job).Error; err != nil {
tx.Rollback()
resp.ERROR(c, err.Error())
return
}
// 如果任务未完成，或者任务失败，则恢复用户算力
if job.Progress != 100 {
err := tx.Model(&model.User{}).Where("id = ?", job.UserId).UpdateColumn("power", gorm.Expr("power + ?", job.Power)).Error
if err != nil {
tx.Rollback()
resp.ERROR(c, err.Error())
return
}
var user model.User
h.DB.Where("id = ?", job.UserId).First(&user)
err = tx.Create(&model.PowerLog{
UserId: user.Id,
Username: user.Username,
Type: types.PowerConsume,
Amount: job.Power,
Balance: user.Power,
Mark: types.PowerAdd,
Model: "dall-e-3",
Remark: fmt.Sprintf("任务失败，退回算力。任务ID：%d，Err: %s", job.Id, job.ErrMsg),
CreatedAt: time.Now(),
}).Error
if err != nil {
tx.Rollback()
resp.ERROR(c, err.Error())
return
}
}
tx.Commit()
// remove image
err := h.uploader.GetUploadHandler().Delete(job.ImgURL)
if err != nil {
@@ -241,10 +273,10 @@ func (h *DallJobHandler) Remove(c *gin.Context) {
// Publish 发布/取消发布图片到画廊显示
func (h *DallJobHandler) Publish(c *gin.Context) {
id := h.GetInt(c, "id", 0)
userId := h.GetInt(c, "user_id", 0)
userId := h.GetLoginUserId(c)
action := h.GetBool(c, "action") // 发布动作，true => 发布，false => 取消分享
res := h.DB.Model(&model.DallJob{Id: uint(id), UserId: uint(userId)}).UpdateColumn("publish", action)
res := h.DB.Model(&model.DallJob{Id: uint(id), UserId: userId}).UpdateColumn("publish", action)
if res.Error != nil {
logger.Error("error with update database：", res.Error)
resp.ERROR(c, "更新数据库失败")

9
api/handler/markmap_handler.go
View File

@@ -101,17 +101,13 @@ func (h *MarkMapHandler) sendMessage(client *types.WsClient, prompt string, mode
return fmt.Errorf("error with query chat model: %v", res.Error)
}
if user.Status == false {
return errors.New("当前用户被禁用")
}
if user.Power < chatModel.Power {
return fmt.Errorf("您当前剩余算力（%d）已不足以支付当前模型算力（%d）！", user.Power, chatModel.Power)
}
messages := make([]interface{}, 0)
messages = append(messages, types.Message{Role: "system", Content: `
你是一位非常优秀的思维导图助手，你会把用户的所有提问都总结成思维导图，然后以 Markdown 格式输出。markdown 只需要输出一级标题，二级标题，三级标题，四级标题，最多输出四级，除此之外不要输出任何其他 markdown 标记。下面是一个合格的例子：
你是一位非常优秀的思维导图助手， 你能帮助用户整理思路，根据用户提供的主题或内容，快速生成结构清晰，有条理的思维导图，然后以 Markdown 格式输出。markdown 只需要输出一级标题，二级标题，三级标题，四级标题，最多输出四级，除此之外不要输出任何其他 markdown 标记。下面是一个合格的例子：
# Geek-AI 助手
## 完整的开源系统
@@ -130,7 +126,7 @@ func (h *MarkMapHandler) sendMessage(client *types.WsClient, prompt string, mode
另外，除此之外不要任何解释性语句。
`})
messages = append(messages, types.Message{Role: "user", Content: prompt})
messages = append(messages, types.Message{Role: "user", Content: fmt.Sprintf("请生成一份有关【%s】一份思维导图，要求结构清晰，有条理", prompt)})
var req = types.ApiRequest{
Model: chatModel.Value,
Stream: true,
@@ -253,5 +249,6 @@ func (h *MarkMapHandler) doRequest(req types.ApiRequest, chatModel model.ChatMod
client = http.DefaultClient
}
request.Header.Set("Authorization", fmt.Sprintf("Bearer %s", apiKey.Value))
logger.Debugf("Sending %s request, API KEY:%s, PROXY: %s, Model: %s", apiKey.ApiURL, apiURL, apiKey.ProxyURL, req.Model)
return client.Do(request)
}

52
api/handler/mj_handler.go
View File

@@ -465,35 +465,37 @@ func (h *MidJourneyHandler) Remove(c *gin.Context) {
return
}
// refund power
err := tx.Model(&model.User{}).Where("id = ?", job.UserId).UpdateColumn("power", gorm.Expr("power + ?", job.Power)).Error
if err != nil {
tx.Rollback()
resp.ERROR(c, err.Error())
return
}
var user model.User
h.DB.Where("id = ?", job.UserId).First(&user)
err = tx.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(),
}).Error
if err != nil {
tx.Rollback()
resp.ERROR(c, err.Error())
return
// 如果任务未完成，或者任务失败，则恢复用户算力
if job.Progress != 100 {
err := tx.Model(&model.User{}).Where("id = ?", job.UserId).UpdateColumn("power", gorm.Expr("power + ?", job.Power)).Error
if err != nil {
tx.Rollback()
resp.ERROR(c, err.Error())
return
}
var user model.User
h.DB.Where("id = ?", job.UserId).First(&user)
err = tx.Create(&model.PowerLog{
UserId: user.Id,
Username: user.Username,
Type: types.PowerConsume,
Amount: job.Power,
Balance: user.Power,
Mark: types.PowerAdd,
Model: "mid-journey",
Remark: fmt.Sprintf("绘画任务失败，退回算力。任务ID：%s，Err: %s", job.TaskId, job.ErrMsg),
CreatedAt: time.Now(),
}).Error
if err != nil {
tx.Rollback()
resp.ERROR(c, err.Error())
return
}
}
tx.Commit()
// remove image
err = h.uploader.GetUploadHandler().Delete(job.ImgURL)
err := h.uploader.GetUploadHandler().Delete(job.ImgURL)
if err != nil {
logger.Error("remove image failed: ", err)
}

46
api/handler/sd_handler.go
View File

@@ -232,7 +232,7 @@ func (h *SdJobHandler) getData(finish bool, userId uint, page int, pageSize int,
session := h.DB.Session(&gorm.Session{})
if finish {
session = session.Where("progress = ?", 100).Order("id DESC")
session = session.Where("progress >= ?", 100).Order("id DESC")
} else {
session = session.Where("progress < ?", 100).Order("id ASC")
}
@@ -278,20 +278,50 @@ func (h *SdJobHandler) getData(finish bool, userId uint, page int, pageSize int,
// Remove remove task image
func (h *SdJobHandler) Remove(c *gin.Context) {
id := h.GetInt(c, "id", 0)
userId := h.GetInt(c, "user_id", 0)
userId := h.GetLoginUserId(c)
var job model.SdJob
if res := h.DB.Where("id = ? AND user_id = ?", id, userId).First(&job); res.Error != nil {
resp.ERROR(c, "记录不存在")
return
}
// remove job recode
res := h.DB.Delete(&model.SdJob{Id: job.Id})
if res.Error != nil {
resp.ERROR(c, res.Error.Error())
// 删除任务
tx := h.DB.Begin()
if err := tx.Delete(&job).Error; err != nil {
tx.Rollback()
resp.ERROR(c, err.Error())
return
}
// 如果任务未完成，或者任务失败，则恢复用户算力
if job.Progress != 100 {
err := tx.Model(&model.User{}).Where("id = ?", job.UserId).UpdateColumn("power", gorm.Expr("power + ?", job.Power)).Error
if err != nil {
tx.Rollback()
resp.ERROR(c, err.Error())
return
}
var user model.User
h.DB.Where("id = ?", job.UserId).First(&user)
err = tx.Create(&model.PowerLog{
UserId: user.Id,
Username: user.Username,
Type: types.PowerConsume,
Amount: job.Power,
Balance: user.Power,
Mark: types.PowerAdd,
Model: "stable-diffusion",
Remark: fmt.Sprintf("任务失败，退回算力。任务ID：%s， Err: %s", job.TaskId, job.ErrMsg),
CreatedAt: time.Now(),
}).Error
if err != nil {
tx.Rollback()
resp.ERROR(c, err.Error())
return
}
}
tx.Commit()
// remove image
err := h.uploader.GetUploadHandler().Delete(job.ImgURL)
if err != nil {
@@ -309,10 +339,10 @@ func (h *SdJobHandler) Remove(c *gin.Context) {
// Publish 发布/取消发布图片到画廊显示
func (h *SdJobHandler) Publish(c *gin.Context) {
id := h.GetInt(c, "id", 0)
userId := h.GetInt(c, "user_id", 0)
userId := h.GetLoginUserId(c)
action := h.GetBool(c, "action") // 发布动作，true => 发布，false => 取消分享
res := h.DB.Model(&model.SdJob{Id: uint(id), UserId: userId}).UpdateColumn("publish", action)
res := h.DB.Model(&model.SdJob{Id: uint(id), UserId: int(userId)}).UpdateColumn("publish", action)
if res.Error != nil {
logger.Error("error with update database：", res.Error)
resp.ERROR(c, "更新数据库失败")

37
api/handler/suno_handler.go
View File

@@ -210,7 +210,42 @@ func (h *SunoHandler) Remove(c *gin.Context) {
return
}
// 删除任务
h.DB.Delete(&job)
tx := h.DB.Begin()
if err := tx.Delete(&job).Error; err != nil {
tx.Rollback()
resp.ERROR(c, err.Error())
return
}
// 如果任务未完成，或者任务失败，则恢复用户算力
if job.Progress != 100 {
err := tx.Model(&model.User{}).Where("id = ?", job.UserId).UpdateColumn("power", gorm.Expr("power + ?", job.Power)).Error
if err != nil {
tx.Rollback()
resp.ERROR(c, err.Error())
return
}
var user model.User
h.DB.Where("id = ?", job.UserId).First(&user)
err = tx.Create(&model.PowerLog{
UserId: user.Id,
Username: user.Username,
Type: types.PowerConsume,
Amount: job.Power,
Balance: user.Power,
Mark: types.PowerAdd,
Model: job.ModelName,
Remark: fmt.Sprintf("Suno 任务失败，退回算力。任务ID：%s，Err:%s", job.TaskId, job.ErrMsg),
CreatedAt: time.Now(),
}).Error
if err != nil {
tx.Rollback()
resp.ERROR(c, err.Error())
return
}
}
tx.Commit()
// 删除文件
_ = h.uploader.GetUploadHandler().Delete(job.CoverURL)
_ = h.uploader.GetUploadHandler().Delete(job.AudioURL)

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

@@ -70,7 +70,7 @@ func (s *Service) Run() {
if err != nil {
logger.Errorf("error with image task: %v", err)
s.db.Model(&model.DallJob{Id: task.JobId}).UpdateColumns(map[string]interface{}{
"progress": -1,
"progress": 101,
"err_msg": err.Error(),
})
s.notifyQueue.RPush(sd.NotifyMessage{UserId: int(task.UserId), JobId: int(task.JobId), Message: sd.Failed})
@@ -148,7 +148,7 @@ func (s *Service) Image(task types.DallTask, sync bool) (string, error) {
Where("enabled", true).
Order("last_used_at ASC").First(&apiKey)
if tx.Error != nil {
return "", fmt.Errorf("no available IMG api key: %v", tx.Error)
return "", fmt.Errorf("no available DALL-E api key: %v", tx.Error)
}
var res imgRes
@@ -225,6 +225,30 @@ func (s *Service) CheckTaskNotify() {
}()
}
func (s *Service) CheckTaskStatus() {
go func() {
logger.Info("Running DALL-E task status checking ...")
for {
var jobs []model.DallJob
res := s.db.Where("progress < ?", 100).Find(&jobs)
if res.Error != nil {
time.Sleep(5 * time.Second)
continue
}
for _, job := range jobs {
// 超时的任务标记为失败
if time.Now().Sub(job.CreatedAt) > time.Minute*10 {
job.Progress = 101
job.ErrMsg = "任务超时"
s.db.Updates(&job)
}
}
time.Sleep(time.Second * 10)
}
}()
}
func (s *Service) DownloadImages() {
go func() {
var items []model.DallJob
@@ -271,44 +295,3 @@ func (s *Service) downloadImage(jobId uint, userId int, orgURL string) (string,
s.notifyQueue.RPush(sd.NotifyMessage{UserId: userId, JobId: int(jobId), Message: sd.Finished})
return imgURL, nil
}
// CheckTaskStatus 检查任务状态，自动删除过期或者失败的任务
func (s *Service) CheckTaskStatus() {
go func() {
logger.Info("Running Stable-Diffusion task status checking ...")
for {
var jobs []model.DallJob
res := s.db.Where("progress < ?", 100).Find(&jobs)
if res.Error != nil {
time.Sleep(5 * time.Second)
continue
}
for _, job := range jobs {
// 5 分钟还没完成的任务直接删除
if time.Now().Sub(job.CreatedAt) > time.Minute*5 || job.Progress == -1 {
s.db.Delete(&job)
var user model.User
s.db.Where("id = ?", job.UserId).First(&user)
// 退回绘图次数
res = s.db.Model(&model.User{}).Where("id = ?", job.UserId).UpdateColumn("power", gorm.Expr("power + ?", job.Power))
if res.Error == nil && res.RowsAffected > 0 {
s.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: "dall-e-3",
Remark: fmt.Sprintf("任务失败，退回算力。任务ID：%d", job.Id),
CreatedAt: time.Now(),
})
}
continue
}
}
time.Sleep(time.Second * 10)
}
}()
}

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

@@ -187,6 +187,14 @@ func (p *ServicePool) SyncTaskProgress() {
}
for _, job := range jobs {
// 5 分钟还没完成的任务标记为失败
if time.Now().Sub(job.CreatedAt) > time.Minute*5 {
job.Progress = 101
job.ErrMsg = "任务超时"
p.db.Updates(&job)
continue
}
if servicePlus := p.getService(job.ChannelId); servicePlus != nil {
_ = servicePlus.Notify(job)
}

26
api/service/sd/pool.go
View File

@@ -109,27 +109,11 @@ func (p *ServicePool) CheckTaskStatus() {
}
for _, job := range jobs {
// 5 分钟还没完成的任务直接删除
if time.Now().Sub(job.CreatedAt) > time.Minute*5 || job.Progress == -1 {
p.db.Delete(&job)
var user model.User
p.db.Where("id = ?", job.UserId).First(&user)
// 退回绘图次数
res = p.db.Model(&model.User{}).Where("id = ?", job.UserId).UpdateColumn("power", gorm.Expr("power + ?", job.Power))
if res.Error == nil && res.RowsAffected > 0 {
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: "stable-diffusion",
Remark: fmt.Sprintf("任务失败，退回算力。任务ID：%s", job.TaskId),
CreatedAt: time.Now(),
})
}
continue
// 5 分钟还没完成的任务标记为失败
if time.Now().Sub(job.CreatedAt) > time.Minute*5 {
job.Progress = 101
job.ErrMsg = "任务超时"
p.db.Updates(&job)
}
}
time.Sleep(time.Second * 5)

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

@@ -87,7 +87,7 @@ func (s *Service) Run() {
logger.Error("绘画任务执行失败：", err.Error())
// update the task progress
s.db.Model(&model.SdJob{Id: uint(task.Id)}).UpdateColumns(map[string]interface{}{
"progress": -1,
"progress": 101,
"err_msg": err.Error(),
})
// 通知前端，任务失败