From 22ae7dd1f3dbccc4918ab25183993b1bba681cf6 Mon Sep 17 00:00:00 2001
From: RockYang <yangjian102621@gmail.com>
Date: Thu, 26 Oct 2023 14:38:06 +0800
Subject: [PATCH] feat: different AI model consuming different amounts of
 use_calls

---
 CHANGELOG.md                            |  7 +++++++
 README.md                               |  2 +-
 api/core/types/chat.go                  |  1 +
 api/handler/chatimpl/azure_handler.go   |  4 +---
 api/handler/chatimpl/baidu_handler.go   |  4 +---
 api/handler/chatimpl/chat_handler.go    | 13 +++++++++++++
 api/handler/chatimpl/chatglm_handler.go |  4 +---
 api/handler/chatimpl/openai_handler.go  |  4 +---
 api/handler/chatimpl/xunfei_handler.go  |  4 +---
 9 files changed, 27 insertions(+), 16 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 0709092f..66bdeaf4 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,5 +1,12 @@
 # 更新日志
 
+## v3.1.7
+1. 功能新增：支持文心4.0 AI 模型
+2. 功能新增：可以在管理后台为用户绑定指定的 AI 模型，如只给某个用户使用 GPT-4 模型
+3. 功能新增：模型新增权重字段，不同的模型每次调用耗费的点数可以设置不同，比如GPT4是GPT3.5的10倍
+4. 功能新增：新增系统配置关闭 AI 模型的函数功能
+5. 功能优化：优化 MidJourney 专业绘画页面图片预览样式
+
 ## v3.1.6
 1. 功能新增：新增AI 绘画照片墙功能页面，供用户查看所有的 AI 绘画作品
 2. 功能新增：新增 AI 角色应用功能页面，用户可以添加自己感兴趣的应用
diff --git a/README.md b/README.md
index 149a41ae..3cfae061 100644
--- a/README.md
+++ b/README.md
@@ -133,7 +133,7 @@ cd docker/mysql
 # 创建 mysql 容器
 docker-compose up -d
 # 导入数据库
-docker exec -i chatgpt-plus-mysql sh -c 'exec mysql -uroot -p12345678' < ../../database/chatgpt_plus-v3.1.6.sql
+docker exec -i chatgpt-plus-mysql sh -c 'exec mysql -uroot -p12345678' < ../../database/chatgpt_plus-v3.1.7.sql
 ```
 
 如果你本地已经安装了 MySQL 服务，那么你只需手动导入数据库即可。
diff --git a/api/core/types/chat.go b/api/core/types/chat.go
index 40987c08..1a107052 100644
--- a/api/core/types/chat.go
+++ b/api/core/types/chat.go
@@ -47,6 +47,7 @@ type ChatModel struct {
 	Id       uint     `json:"id"`
 	Platform Platform `json:"platform"`
 	Value    string   `json:"value"`
+	Weight   int      `json:"weight"`
 }
 
 type ApiError struct {
diff --git a/api/handler/chatimpl/azure_handler.go b/api/handler/chatimpl/azure_handler.go
index 49c21276..c3b44f7e 100644
--- a/api/handler/chatimpl/azure_handler.go
+++ b/api/handler/chatimpl/azure_handler.go
@@ -168,9 +168,7 @@ func (h *ChatHandler) sendAzureMessage(
 		// 消息发送成功
 		if len(contents) > 0 {
 			// 更新用户的对话次数
-			if userVo.ChatConfig.ApiKeys[session.Model.Platform] == "" {
-				h.db.Model(&model.User{}).Where("id = ?", userVo.Id).UpdateColumn("calls", gorm.Expr("calls - ?", 1))
-			}
+			h.subUserCalls(userVo, session)
 
 			if message.Role == "" {
 				message.Role = "assistant"
diff --git a/api/handler/chatimpl/baidu_handler.go b/api/handler/chatimpl/baidu_handler.go
index fc545dce..b0b7e7cf 100644
--- a/api/handler/chatimpl/baidu_handler.go
+++ b/api/handler/chatimpl/baidu_handler.go
@@ -124,9 +124,7 @@ func (h *ChatHandler) sendBaiduMessage(
 		// 消息发送成功
 		if len(contents) > 0 {
 			// 更新用户的对话次数
-			if userVo.ChatConfig.ApiKeys[session.Model.Platform] == "" {
-				h.db.Model(&model.User{}).Where("id = ?", userVo.Id).UpdateColumn("calls", gorm.Expr("calls - ?", 1))
-			}
+			h.subUserCalls(userVo, session)
 
 			if message.Role == "" {
 				message.Role = "assistant"
diff --git a/api/handler/chatimpl/chat_handler.go b/api/handler/chatimpl/chat_handler.go
index 83493d4e..db13b9a0 100644
--- a/api/handler/chatimpl/chat_handler.go
+++ b/api/handler/chatimpl/chat_handler.go
@@ -103,6 +103,7 @@ func (h *ChatHandler) ChatHandle(c *gin.Context) {
 	session.Model = types.ChatModel{
 		Id:       chatModel.Id,
 		Value:    chatModel.Value,
+		Weight:   chatModel.Weight,
 		Platform: types.Platform(chatModel.Platform)}
 	logger.Infof("New websocket connected, IP: %s, Username: %s", c.ClientIP(), session.Username)
 	var chatRole model.ChatRole
@@ -463,3 +464,15 @@ func (h *ChatHandler) doRequest(ctx context.Context, req types.ApiRequest, platf
 	}
 	return client.Do(request)
 }
+
+// 扣减用户的对话次数
+func (h *ChatHandler) subUserCalls(userVo vo.User, session *types.ChatSession) {
+	// 仅当用户没有导入自己的 API KEY 时才进行扣减
+	if userVo.ChatConfig.ApiKeys[session.Model.Platform] == "" {
+		num := 1
+		if session.Model.Weight > 0 {
+			num = session.Model.Weight
+		}
+		h.db.Model(&model.User{}).Where("id = ?", userVo.Id).UpdateColumn("calls", gorm.Expr("calls - ?", num))
+	}
+}
diff --git a/api/handler/chatimpl/chatglm_handler.go b/api/handler/chatimpl/chatglm_handler.go
index c27ffa62..25df5b9e 100644
--- a/api/handler/chatimpl/chatglm_handler.go
+++ b/api/handler/chatimpl/chatglm_handler.go
@@ -104,9 +104,7 @@ func (h *ChatHandler) sendChatGLMMessage(
 		// 消息发送成功
 		if len(contents) > 0 {
 			// 更新用户的对话次数
-			if userVo.ChatConfig.ApiKeys[session.Model.Platform] == "" {
-				h.db.Model(&model.User{}).Where("id = ?", userVo.Id).UpdateColumn("calls", gorm.Expr("calls - ?", 1))
-			}
+			h.subUserCalls(userVo, session)
 
 			if message.Role == "" {
 				message.Role = "assistant"
diff --git a/api/handler/chatimpl/openai_handler.go b/api/handler/chatimpl/openai_handler.go
index f7756a97..478bba19 100644
--- a/api/handler/chatimpl/openai_handler.go
+++ b/api/handler/chatimpl/openai_handler.go
@@ -167,9 +167,7 @@ func (h *ChatHandler) sendOpenAiMessage(
 		// 消息发送成功
 		if len(contents) > 0 {
 			// 更新用户的对话次数
-			if userVo.ChatConfig.ApiKeys[session.Model.Platform] == "" {
-				h.db.Model(&model.User{}).Where("id = ?", userVo.Id).UpdateColumn("calls", gorm.Expr("calls - ?", 1))
-			}
+			h.subUserCalls(userVo, session)
 
 			if message.Role == "" {
 				message.Role = "assistant"
diff --git a/api/handler/chatimpl/xunfei_handler.go b/api/handler/chatimpl/xunfei_handler.go
index 5f1427e5..19710568 100644
--- a/api/handler/chatimpl/xunfei_handler.go
+++ b/api/handler/chatimpl/xunfei_handler.go
@@ -167,9 +167,7 @@ func (h *ChatHandler) sendXunFeiMessage(
 	// 消息发送成功
 	if len(contents) > 0 {
 		// 更新用户的对话次数
-		if userVo.ChatConfig.ApiKeys[session.Model.Platform] == "" {
-			h.db.Model(&model.User{}).Where("id = ?", userVo.Id).UpdateColumn("calls", gorm.Expr("calls - ?", 1))
-		}
+		h.subUserCalls(userVo, session)
 
 		if message.Role == "" {
 			message.Role = "assistant"