From bbc8fe2b400f0dde4a7cbc749983767f1b933d4c Mon Sep 17 00:00:00 2001
From: RockYang <yangjian102621@gmail.com>
Date: Fri, 10 May 2024 11:18:37 +0800
Subject: [PATCH] fixed bug for dalle3 task not decrease power

---
 CHANGELOG.md                   |  1 +
 api/handler/dalle_handler.go   |  3 +--
 api/handler/markmap_handler.go | 19 +++++++++++++++++++
 3 files changed, 21 insertions(+), 2 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 8b9454cc..faabe18d 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -4,6 +4,7 @@
 
 * Bug修复：修复PC端画廊页面的瀑布流组件样式错乱问题
 * 功能新增：给思维导图增加 ToolBar，实现思维导图的放大缩小和定位
+* Bug修复：修复思维导图不扣费的Bug
 
 ## v4.0.5
 
diff --git a/api/handler/dalle_handler.go b/api/handler/dalle_handler.go
index ede79696..a62fd712 100644
--- a/api/handler/dalle_handler.go
+++ b/api/handler/dalle_handler.go
@@ -92,8 +92,7 @@ func (h *DallJobHandler) preCheck(c *gin.Context) bool {
 		resp.NotAuth(c)
 		return false
 	}
-
-	if user.Power < h.App.SysConfig.SdPower {
+	if user.Power < h.App.SysConfig.DallPower {
 		resp.ERROR(c, "当前用户剩余算力不足以完成本次绘画！")
 		return false
 	}
diff --git a/api/handler/markmap_handler.go b/api/handler/markmap_handler.go
index 248e5802..e74f5036 100644
--- a/api/handler/markmap_handler.go
+++ b/api/handler/markmap_handler.go
@@ -206,6 +206,25 @@ func (h *MarkMapHandler) sendMessage(client *types.WsClient, prompt string, mode
 		}
 	}
 
+	// 扣减算力
+	res = h.DB.Model(&model.User{}).Where("id", userId).UpdateColumn("power", gorm.Expr("power - ?", chatModel.Power))
+	if res.Error == nil {
+		// 记录算力消费日志
+		var u model.User
+		h.DB.Where("id", userId).First(&u)
+		h.DB.Create(&model.PowerLog{
+			UserId:    u.Id,
+			Username:  u.Username,
+			Type:      types.PowerConsume,
+			Amount:    chatModel.Power,
+			Mark:      types.PowerSub,
+			Balance:   u.Power,
+			Model:     chatModel.Value,
+			Remark:    fmt.Sprintf("AI绘制思维导图，模型名称：%s, ", chatModel.Value),
+			CreatedAt: time.Now(),
+		})
+	}
+
 	return nil
 }