From 7060edb3e5eabb0971fd72a22a1fc63a76d43c2f Mon Sep 17 00:00:00 2001
From: RockYang <yangjian102621@gmail.com>
Date: Thu, 21 Mar 2024 15:55:39 +0800
Subject: [PATCH] feat: save prompt in power log for dalle-3

---
 api/handler/chatimpl/chat_handler.go |  4 +++-
 api/handler/function_handler.go      |  8 +++++---
 api/utils/strings.go                 | 16 ++++++++++------
 3 files changed, 18 insertions(+), 10 deletions(-)

diff --git a/api/handler/chatimpl/chat_handler.go b/api/handler/chatimpl/chat_handler.go
index 11e10ebe..04066167 100644
--- a/api/handler/chatimpl/chat_handler.go
+++ b/api/handler/chatimpl/chat_handler.go
@@ -528,13 +528,15 @@ func (h *ChatHandler) subUserPower(userVo vo.User, session *types.ChatSession, p
 	res := h.DB.Model(&model.User{}).Where("id = ?", userVo.Id).UpdateColumn("power", gorm.Expr("power - ?", power))
 	if res.Error == nil {
 		// 记录算力消费日志
+		var u model.User
+		h.DB.Where("id", userVo.Id).First(&u)
 		h.DB.Create(&model.PowerLog{
 			UserId:    userVo.Id,
 			Username:  userVo.Username,
 			Type:      types.PowerConsume,
 			Amount:    power,
 			Mark:      types.PowerSub,
-			Balance:   userVo.Power - power,
+			Balance:   u.Power,
 			Model:     session.Model.Value,
 			Remark:    fmt.Sprintf("模型名称：%s, 提问长度：%d，回复长度：%d", session.Model.Name, promptTokens, replyTokens),
 			CreatedAt: time.Now(),
diff --git a/api/handler/function_handler.go b/api/handler/function_handler.go
index a7ef6af3..cecbf057 100644
--- a/api/handler/function_handler.go
+++ b/api/handler/function_handler.go
@@ -252,18 +252,20 @@ func (h *FunctionHandler) Dall3(c *gin.Context) {
 
 	content := fmt.Sprintf("下面是根据您的描述创作的图片，它描绘了 【%s】 的场景。 \n\n![](%s)\n", prompt, imgURL)
 	// 更新用户算力
-	tx = h.DB.Model(&model.User{}).Where("id = ?", user.Id).UpdateColumn("power", gorm.Expr("power - ?", h.App.SysConfig.DallPower))
+	tx = h.DB.Model(&model.User{}).Where("id", user.Id).UpdateColumn("power", gorm.Expr("power - ?", h.App.SysConfig.DallPower))
 	// 记录算力变化日志
 	if tx.Error == nil && tx.RowsAffected > 0 {
+		var u model.User
+		h.DB.Where("id", user.Id).First(&u)
 		h.DB.Create(&model.PowerLog{
 			UserId:    user.Id,
 			Username:  user.Username,
 			Type:      types.PowerConsume,
 			Amount:    h.App.SysConfig.DallPower,
-			Balance:   user.Power - h.App.SysConfig.DallPower,
+			Balance:   u.Power,
 			Mark:      types.PowerSub,
 			Model:     "dall-e-3",
-			Remark:    "",
+			Remark:    fmt.Sprintf("绘画提示词：%s", utils.CutWords(prompt, 10)),
 			CreatedAt: time.Now(),
 		})
 	}
diff --git a/api/utils/strings.go b/api/utils/strings.go
index ccccbc97..db4b194d 100644
--- a/api/utils/strings.go
+++ b/api/utils/strings.go
@@ -4,7 +4,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"math/rand"
-	"strconv"
+	"strings"
 	"time"
 
 	"golang.org/x/crypto/sha3"
@@ -94,10 +94,14 @@ func InterfaceToString(value interface{}) string {
 	return JsonEncode(value)
 }
 
-func Str2Float(str string) float64 {
-	num, err := strconv.ParseFloat(str, 64)
-	if err != nil {
-		return 0
+func CutWords(str string, num int) string {
+	// 按空格分割字符串为单词切片
+	words := strings.Fields(str)
+
+	// 如果单词数量超过指定数量，则裁剪单词；否则保持原样
+	if len(words) > num {
+		return strings.Join(words[:num], " ") + " ..."
+	} else {
+		return str
 	}
-	return num
 }