restore use power when removed not finish jobs

api/handler/chatimpl/chat_handler.go

@@ -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())

api/handler/dalle_handler.go

@@ -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, "更新数据库失败")

api/handler/markmap_handler.go

@@ -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)
 }

api/handler/mj_handler.go

@@ -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)
 	}

api/handler/sd_handler.go

@@ -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, "更新数据库失败")

api/handler/suno_handler.go

@@ -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)