From f747bf5ead72c048a6e28d5a5857b04c12c11cb6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E6=9D=8E=E5=AF=BB=E6=AC=A2?= <lxh@cxh.cn>
Date: Fri, 5 Jul 2024 09:44:07 +0800
Subject: [PATCH] =?UTF-8?q?:sparkles:=20=E6=96=B0=E5=A2=9E=E8=AE=B0?=
 =?UTF-8?q?=E5=BD=95=E7=BE=A4=E6=88=96=E8=80=85=E5=A5=BD=E5=8F=8B=E5=B7=B2?=
 =?UTF-8?q?=E4=BD=BF=E7=94=A8=E7=9A=84tokens=E4=BF=A1=E6=81=AF?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 model/entity/friend.go |  1 +
 plugin/plugins/ai.go   |  3 +++
 service/friend.go      | 14 ++++++++++++++
 3 files changed, 18 insertions(+)

diff --git a/model/entity/friend.go b/model/entity/friend.go
index 884c27d5..26796c7a 100644
--- a/model/entity/friend.go
+++ b/model/entity/friend.go
@@ -22,6 +22,7 @@ type Friend struct {
 	EnableNews     bool      `json:"enableNews" gorm:"type:tinyint(1) default 0 not null"`     // 是否启用新闻
 	ClearMember    int       `json:"clearMember"`                                              // 清理成员配置（多少天未活跃的）
 	IsOk           bool      `json:"isOk" gorm:"type:tinyint(1) default 0 not null"`           // 是否正常
+	UsedTokens     int       `json:"usedTokens"`                                               // 已使用的AI Token数量
 }
 
 func (Friend) TableName() string {
diff --git a/plugin/plugins/ai.go b/plugin/plugins/ai.go
index 1edbb02f..edfcc5c8 100644
--- a/plugin/plugins/ai.go
+++ b/plugin/plugins/ai.go
@@ -142,6 +142,9 @@ func AI(m *plugin.MessageContext) {
 		return
 	}
 
+	// 异步更新一下已使用的AI次数
+	go service.UpdateUsedAiTokens(m.FromUser, resp.Usage.TotalTokens)
+
 	// 保存一下AI 返回的消息，消息 Id 使用传入 Id 的负数
 	var replyMessage entity.Message
 	replyMessage.MsgId = -m.MsgId
diff --git a/service/friend.go b/service/friend.go
index 1729f341..95be7af7 100644
--- a/service/friend.go
+++ b/service/friend.go
@@ -4,6 +4,7 @@ import (
 	"go-wechat/client"
 	"go-wechat/model/entity"
 	"go-wechat/model/vo"
+	"gorm.io/gorm"
 	"log"
 	"strings"
 )
@@ -119,3 +120,16 @@ func updateLastActive(msg entity.Message) {
 		log.Printf("更新群或者好友活跃时间失败, 错误信息: %v", err)
 	}
 }
+
+// UpdateUsedAiTokens
+// @description: 更新已使用的AI次数
+// @param wxId 微信好友或者群聊Id
+// @param tokens 新增的tokens额度
+func UpdateUsedAiTokens(wxId string, tokens int) {
+	err := client.MySQL.Model(&entity.Friend{}).
+		Where("wxid = ?", wxId).
+		Update("`used_tokens`", gorm.Expr(" `used_tokens` + ?", tokens)).Error
+	if err != nil {
+		log.Printf("更新AI使用次数失败, 错误信息: %v", err)
+	}
+}