feat: add authorization for MidJourney function calls

api/core/types/config.go

@@ -93,6 +93,7 @@ type SystemConfig struct {
 	AdminTitle        string   `json:"admin_title"`
 	Models            []string `json:"models"`
 	UserInitCalls     int      `json:"user_init_calls"` // 新用户注册默认总送多少次调用
+	InitImgCalls      int      `json:"init_img_calls"`
 	EnabledRegister   bool     `json:"enabled_register"`
 	EnabledMsgService bool     `json:"enabled_msg_service"`
 }

api/handler/admin/user_handler.go

@@ -73,6 +73,7 @@ func (h *UserHandler) Save(c *gin.Context) {
 		Mobile      string   `json:"mobile"`
 		Nickname    string   `json:"nickname"`
 		Calls       int      `json:"calls"`
+		ImgCalls    int      `json:"img_calls"`
 		ChatRoles   []string `json:"chat_roles"`
 		ExpiredTime string   `json:"expired_time"`
 		Status      bool     `json:"status"`
@@ -91,6 +92,7 @@ func (h *UserHandler) Save(c *gin.Context) {
 			"nickname":        data.Nickname,
 			"mobile":          data.Mobile,
 			"calls":           data.Calls,
+			"img_calls":       data.ImgCalls,
 			"status":          data.Status,
 			"chat_roles_json": utils.JsonEncode(data.ChatRoles),
 			"expired_time":    utils.Str2stamp(data.ExpiredTime),

api/handler/chat_handler.go

@@ -326,10 +326,15 @@ func (h *ChatHandler) sendMessage(ctx context.Context, session types.ChatSession
 			} // end for
 
 			if functionCall { // 调用函数完成任务
-				logger.Info("函数名称：", functionName)
 				var params map[string]interface{}
 				_ = utils.JsonDecode(strings.Join(arguments, ""), &params)
-				logger.Info("函数参数：", params)
+				logger.Debugf("函数名称: %s, 函数参数：%s", functionName, params)
+
+				// for creating image, check if the user's img_calls > 0
+				if functionName == types.FuncMidJourney && userVo.ImgCalls <= 0 {
+					utils.ReplyMessage(ws, "**当前用户剩余绘图次数已用尽，请扫描下面二维码联系管理员！**")
+					utils.ReplyMessage(ws, "![](/images/wx.png)")
+				} else {
 					f := h.App.Functions[functionName]
 					data, err := f.Invoke(params)
 					if err != nil {
@@ -343,7 +348,7 @@ func (h *ChatHandler) sendMessage(ctx context.Context, session types.ChatSession
 						content := data
 						if functionName == types.FuncMidJourney {
 							key := utils.Sha256(data)
-						//logger.Info(data, ",", key)
+							logger.Debug(data, ",", key)
 							// add task for MidJourney
 							h.App.MjTaskClients.Put(key, ws)
 							task := types.MjTask{
@@ -357,6 +362,9 @@ func (h *ChatHandler) sendMessage(ctx context.Context, session types.ChatSession
 								logger.Error("error with store MidJourney task: ", err)
 							}
 							content = fmt.Sprintf("绘画提示词：%s 已推送任务到 MidJourney 机器人，请耐心等待任务执行...", data)
+
+							// update user's img_calls
+							h.db.Model(&user).UpdateColumn("img_calls", gorm.Expr("img_calls - ?", 1))
 						}
 
 						utils.ReplyChunkMessage(ws, types.WsMessage{
@@ -366,15 +374,13 @@ func (h *ChatHandler) sendMessage(ctx context.Context, session types.ChatSession
 						contents = append(contents, content)
 					}
 				}
+			}
 
 			// 消息发送成功
 			if len(contents) > 0 {
 				// 更新用户的对话次数
 				if userVo.ChatConfig.ApiKey == "" { // 如果用户使用的是自己绑定的 API KEY 则不扣减对话次数
-					res := h.db.Model(&user).UpdateColumn("calls", gorm.Expr("calls - ?", 1))
+					h.db.Model(&user).UpdateColumn("calls", gorm.Expr("calls - ?", 1))
-					if res.Error != nil {
-						return res.Error
-					}
 				}
 
 				if message.Role == "" {

api/handler/mj_handler.go

@@ -109,6 +109,7 @@ func (h *MidJourneyHandler) Notify(c *gin.Context) {
 		job.UserId = task.UserId
 		job.ChatId = task.ChatId
 		job.MessageId = data.MessageId
+		job.ReferenceId = data.ReferenceId
 		job.Content = data.Content
 		job.Prompt = data.Prompt
 		job.Image = utils.JsonEncode(data.Image)

api/handler/user_handler.go

@@ -109,6 +109,7 @@ func (h *UserHandler) Register(c *gin.Context) {
 			ApiKey:        "",
 		}),
 		Calls:    h.App.SysConfig.UserInitCalls,
+		ImgCalls: h.App.SysConfig.InitImgCalls,
 	}
 	res = h.db.Create(&user)
 	if res.Error != nil {

api/store/model/mj_job.go

@@ -7,6 +7,7 @@ type MidJourneyJob struct {
 	UserId      uint
 	ChatId      string
 	MessageId   string
+	ReferenceId string
 	Hash        string
 	Content     string
 	Prompt      string

api/store/model/user.go

@@ -10,6 +10,7 @@ type User struct {
 	Salt        string // 密码盐
 	Tokens      int64  // 剩余tokens
 	Calls       int    // 剩余对话次数
+	ImgCalls    int    // 剩余绘图次数
 	ChatConfig  string `gorm:"column:chat_config_json"` // 聊天配置 json
 	ChatRoles   string `gorm:"column:chat_roles_json"`  // 聊天角色
 	ExpiredTime int64  // 账户到期时间

api/store/vo/user.go

@@ -11,6 +11,7 @@ type User struct {
 	Salt        string           `json:"salt"`   // 密码盐
 	Tokens      int64            `json:"tokens"` // 剩余tokens
 	Calls       int              `json:"calls"`  // 剩余对话次数
+	ImgCalls    int              `json:"img_calls"`
 	ChatConfig  types.ChatConfig `json:"chat_config"`   // 聊天配置
 	ChatRoles   []string         `json:"chat_roles"`    // 聊天角色集合
 	ExpiredTime int64            `json:"expired_time"`  // 账户到期时间

database/update-v3.0.7.sql

@@ -31,4 +31,7 @@ ALTER TABLE `chatgpt_mj_jobs`
 --
 ALTER TABLE `chatgpt_mj_jobs`
     MODIFY `id` int NOT NULL AUTO_INCREMENT;
-COMMIT;
+
+ALTER TABLE `chatgpt_mj_jobs` ADD `reference_id` CHAR(40) NULL DEFAULT NULL COMMENT '引用消息 ID' AFTER `message_id`;
+
+ALTER TABLE `chatgpt_users` ADD `img_calls` INT NOT NULL DEFAULT '0' COMMENT '剩余绘图次数' AFTER `calls`;

web/src/views/admin/RewardList.vue

@@ -6,6 +6,7 @@
         <el-table-column prop="username" label="用户名"/>
         <el-table-column prop="tx_id" label="转账单号"/>
         <el-table-column prop="amount" label="转账金额"/>
+        <el-table-column prop="remark" label="备注"/>
 
         <el-table-column label="转账时间">
           <template #default="scope">
@@ -27,11 +28,10 @@
 </template>
 
 <script setup>
-import {reactive, ref} from "vue";
+import {ref} from "vue";
-import {httpGet, httpPost} from "@/utils/http";
+import {httpGet} from "@/utils/http";
 import {ElMessage} from "element-plus";
-import {dateFormat, disabledDate, removeArrayItem} from "@/utils/libs";
+import {dateFormat} from "@/utils/libs";
-import {Plus} from "@element-plus/icons-vue";
 
 // 变量定义
 const items = ref([])

web/src/views/admin/SysConfig.vue

@@ -9,9 +9,12 @@
         <el-form-item label="控制台标题" prop="admin_title">
           <el-input v-model="system['admin_title']"/>
         </el-form-item>
-        <el-form-item label="注册赠送次数" prop="init_calls">
+        <el-form-item label="赠送对话次数" prop="init_calls">
           <el-input v-model.number="system['user_init_calls']" placeholder="新用户注册赠送对话次数"/>
         </el-form-item>
+        <el-form-item label="赠送绘图次数" prop="init_calls">
+          <el-input v-model.number="system['init_img_calls']" placeholder="新用户注册赠送绘图次数"/>
+        </el-form-item>
         <el-form-item label="短信验证服务" prop="enabled_msg_service">
           <el-switch v-model="system['enabled_msg_service']"/>
         </el-form-item>

web/src/views/admin/UserList.vue

@@ -82,6 +82,9 @@
         <el-form-item label="提问次数：" prop="calls">
           <el-input v-model.number="user.calls" autocomplete="off" placeholder="0"/>
         </el-form-item>
+        <el-form-item label="绘图次数：" prop="img_calls">
+          <el-input v-model.number="user['img_calls']" autocomplete="off" placeholder="0"/>
+        </el-form-item>
 
         <el-form-item label="有效期：" prop="expired_time">
           <el-date-picker