From 49646a79e565c26b231ca308f680d9fc83e126c5 Mon Sep 17 00:00:00 2001
From: RockYang <yangjian102621@gmail.com>
Date: Sun, 16 Jul 2023 11:55:45 +0800
Subject: [PATCH] =?UTF-8?q?fix:=20=E7=94=A8=E6=88=B7=E4=BD=BF=E7=94=A8?=
 =?UTF-8?q?=E8=87=AA=E5=B7=B1=E7=BB=91=E5=AE=9A=E7=9A=84=20API=20=E8=B0=83?=
 =?UTF-8?q?=E7=94=A8=EF=BC=8C=E5=88=99=E4=B8=8D=E8=AE=A1=E7=AE=97=20calls?=
 =?UTF-8?q?=20=E5=92=8C=20token=20=E6=B6=88=E8=80=97?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 api/handler/chat_handler.go | 16 ++++++++++------
 1 file changed, 10 insertions(+), 6 deletions(-)

diff --git a/api/handler/chat_handler.go b/api/handler/chat_handler.go
index 0101a66e..3e683a9a 100644
--- a/api/handler/chat_handler.go
+++ b/api/handler/chat_handler.go
@@ -151,7 +151,7 @@ func (h *ChatHandler) sendMessage(ctx context.Context, session types.ChatSession
 		return nil
 	}
 
-	if userVo.Calls <= 0 {
+	if userVo.Calls <= 0 && userVo.ChatConfig.ApiKey == "" {
 		replyMessage(ws, "您的对话次数已经用尽，请联系管理员充值！")
 		replyMessage(ws, "![](/images/wx.png)")
 		return nil
@@ -327,9 +327,11 @@ func (h *ChatHandler) sendMessage(ctx context.Context, session types.ChatSession
 			// 消息发送成功
 			if len(contents) > 0 {
 				// 更新用户的对话次数
-				res := h.db.Model(&user).UpdateColumn("calls", gorm.Expr("calls - ?", 1))
-				if res.Error != nil {
-					return res.Error
+				if userVo.ChatConfig.ApiKey == "" { // 如果用户使用的是自己绑定的 API KEY 则不扣减对话次数
+					res := h.db.Model(&user).UpdateColumn("calls", gorm.Expr("calls - ?", 1))
+					if res.Error != nil {
+						return res.Error
+					}
 				}
 
 				if message.Role == "" {
@@ -408,8 +410,10 @@ func (h *ChatHandler) sendMessage(ctx context.Context, session types.ChatSession
 						totalTokens += getTotalTokens(req)
 					}
 					//replyChunkMessage(ws, types.WsMessage{Type: types.WsMiddle, Content: fmt.Sprintf("\n\n `本轮对话共消耗 Token 数量: %d`", totalTokens+11)})
-					h.db.Model(&user).UpdateColumn("tokens", gorm.Expr("tokens + ?",
-						totalTokens))
+					if userVo.ChatConfig.ApiKey != "" { // 调用自己的 API KEY 不计算 token 消耗
+						h.db.Model(&user).UpdateColumn("tokens", gorm.Expr("tokens + ?",
+							totalTokens))
+					}
 				}
 
 				// 保存当前会话