diff --git a/api/core/types/config.go b/api/core/types/config.go index f4b87be6..bf0e4c18 100644 --- a/api/core/types/config.go +++ b/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"` } diff --git a/api/handler/admin/user_handler.go b/api/handler/admin/user_handler.go index d491d83a..2312c5b7 100644 --- a/api/handler/admin/user_handler.go +++ b/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), diff --git a/api/handler/chat_handler.go b/api/handler/chat_handler.go index e8d3c18a..94a44061 100644 --- a/api/handler/chat_handler.go +++ b/api/handler/chat_handler.go @@ -326,44 +326,53 @@ 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, ""), ¶ms) - logger.Info("函数参数：", params) - f := h.App.Functions[functionName] - data, err := f.Invoke(params) - if err != nil { - msg := "调用函数出错：" + err.Error() - utils.ReplyChunkMessage(ws, types.WsMessage{ - Type: types.WsMiddle, - Content: msg, - }) - contents = append(contents, msg) - } else { - content := data - if functionName == types.FuncMidJourney { - key := utils.Sha256(data) - //logger.Info(data, ",", key) - // add task for MidJourney - h.App.MjTaskClients.Put(key, ws) - task := types.MjTask{ - UserId: userVo.Id, - RoleId: role.Id, - Icon: "/images/avatar/mid_journey.png", - ChatId: session.ChatId, - } - err := h.leveldb.Put(types.TaskStorePrefix+key, task) - if err != nil { - logger.Error("error with store MidJourney task: ", err) - } - content = fmt.Sprintf("绘画提示词：%s 已推送任务到 MidJourney 机器人，请耐心等待任务执行...", data) - } + logger.Debugf("函数名称: %s, 函数参数：%s", functionName, params) - utils.ReplyChunkMessage(ws, types.WsMessage{ - Type: types.WsMiddle, - Content: content, - }) - contents = append(contents, content) + // 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 { + msg := "调用函数出错：" + err.Error() + utils.ReplyChunkMessage(ws, types.WsMessage{ + Type: types.WsMiddle, + Content: msg, + }) + contents = append(contents, msg) + } else { + content := data + if functionName == types.FuncMidJourney { + key := utils.Sha256(data) + logger.Debug(data, ",", key) + // add task for MidJourney + h.App.MjTaskClients.Put(key, ws) + task := types.MjTask{ + UserId: userVo.Id, + RoleId: role.Id, + Icon: "/images/avatar/mid_journey.png", + ChatId: session.ChatId, + } + err := h.leveldb.Put(types.TaskStorePrefix+key, task) + if err != nil { + 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{ + Type: types.WsMiddle, + Content: content, + }) + contents = append(contents, content) + } } } @@ -371,10 +380,7 @@ func (h *ChatHandler) sendMessage(ctx context.Context, session types.ChatSession if len(contents) > 0 { // 更新用户的对话次数 if userVo.ChatConfig.ApiKey == "" { // 如果用户使用的是自己绑定的 API KEY 则不扣减对话次数 - res := h.db.Model(&user).UpdateColumn("calls", gorm.Expr("calls - ?", 1)) - if res.Error != nil { - return res.Error - } + h.db.Model(&user).UpdateColumn("calls", gorm.Expr("calls - ?", 1)) } if message.Role == "" { diff --git a/api/handler/mj_handler.go b/api/handler/mj_handler.go index aad4309b..ec55da48 100644 --- a/api/handler/mj_handler.go +++ b/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) diff --git a/api/handler/user_handler.go b/api/handler/user_handler.go index 245e9bd1..e25e08cd 100644 --- a/api/handler/user_handler.go +++ b/api/handler/user_handler.go @@ -108,7 +108,8 @@ func (h *UserHandler) Register(c *gin.Context) { Model: h.App.ChatConfig.Model, ApiKey: "", }), - Calls: h.App.SysConfig.UserInitCalls, + Calls: h.App.SysConfig.UserInitCalls, + ImgCalls: h.App.SysConfig.InitImgCalls, } res = h.db.Create(&user) if res.Error != nil { diff --git a/api/store/model/mj_job.go b/api/store/model/mj_job.go index 74f09563..3162273a 100644 --- a/api/store/model/mj_job.go +++ b/api/store/model/mj_job.go @@ -3,15 +3,16 @@ package model import "time" type MidJourneyJob struct { - Id uint `gorm:"primarykey;column:id"` - UserId uint - ChatId string - MessageId string - Hash string - Content string - Prompt string - Image string - CreatedAt time.Time + Id uint `gorm:"primarykey;column:id"` + UserId uint + ChatId string + MessageId string + ReferenceId string + Hash string + Content string + Prompt string + Image string + CreatedAt time.Time } func (MidJourneyJob) TableName() string { diff --git a/api/store/model/user.go b/api/store/model/user.go index 2e2e9684..76644685 100644 --- a/api/store/model/user.go +++ b/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 // 账户到期时间 diff --git a/api/store/vo/user.go b/api/store/vo/user.go index df66132e..37cd1626 100644 --- a/api/store/vo/user.go +++ b/api/store/vo/user.go @@ -8,9 +8,10 @@ type User struct { Mobile string `json:"mobile"` Nickname string `json:"nickname"` Avatar string `json:"avatar"` - Salt string `json:"salt"` // 密码盐 - Tokens int64 `json:"tokens"` // 剩余tokens - Calls int `json:"calls"` // 剩余对话次数 + 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"` // 账户到期时间 diff --git a/database/update-v3.0.7.sql b/database/update-v3.0.7.sql index 5031672a..fc7664f2 100644 --- a/database/update-v3.0.7.sql +++ b/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`; \ No newline at end of file diff --git a/web/src/views/admin/RewardList.vue b/web/src/views/admin/RewardList.vue index cd6b725c..fc13ece2 100644 --- a/web/src/views/admin/RewardList.vue +++ b/web/src/views/admin/RewardList.vue @@ -6,6 +6,7 @@ +