From 2db58ce66c510c191b16b52d934be05b00905f4e Mon Sep 17 00:00:00 2001 From: dean <1006393151@qq.com> Date: Sat, 17 Aug 2024 15:44:20 +0800 Subject: [PATCH] feat: MESSAGE_MAX_TOKEN --- common/constants.go | 2 +- controller/chat.go | 32 ++++++++++++++++---------------- discord/discord.go | 40 +++++++++++++++++++++++++++++++--------- model/bot.go | 9 +++++---- 4 files changed, 53 insertions(+), 30 deletions(-) diff --git a/common/constants.go b/common/constants.go index 8cb3f80b..303529c5 100644 --- a/common/constants.go +++ b/common/constants.go @@ -1,6 +1,6 @@ package common -var Version = "v4.5.10" // this hard coding will be replaced automatically when building, no need to manually change +var Version = "v4.6.0" // this hard coding will be replaced automatically when building, no need to manually change const ( RequestIdKey = "X-Request-Id" diff --git a/controller/chat.go b/controller/chat.go index 46433f1e..2647717a 100644 --- a/controller/chat.go +++ b/controller/chat.go @@ -166,7 +166,7 @@ func ChatForOpenAI(c *gin.Context) { return } - sendChannelId, calledCozeBotId, isNewChannel, err := getSendChannelIdAndCozeBotId(c, request.ChannelId, request.Model, true) + sendChannelId, calledCozeBotId, maxToken, isNewChannel, err := getSendChannelIdAndCozeBotId(c, request.ChannelId, request.Model, true) if err != nil { response := model.OpenAIErrorResponse{ @@ -275,7 +275,7 @@ loop: content = string(jsonData) } - sentMsg, userAuth, err := discord.SendMessage(c, sendChannelId, calledCozeBotId, content) + sentMsg, userAuth, err := discord.SendMessage(c, sendChannelId, calledCozeBotId, content, maxToken) if err != nil { c.JSON(http.StatusInternalServerError, model.OpenAIErrorResponse{ OpenAIError: model.OpenAIError{ @@ -517,7 +517,7 @@ func ImagesForOpenAI(c *gin.Context) { return } - sendChannelId, calledCozeBotId, isNewChannel, err := getSendChannelIdAndCozeBotId(c, request.ChannelId, request.Model, true) + sendChannelId, calledCozeBotId, maxToken, isNewChannel, err := getSendChannelIdAndCozeBotId(c, request.ChannelId, request.Model, true) if err != nil { common.LogError(c.Request.Context(), err.Error()) c.JSON(http.StatusInternalServerError, model.OpenAIErrorResponse{ @@ -547,7 +547,7 @@ func ImagesForOpenAI(c *gin.Context) { }() } - sentMsg, userAuth, err := discord.SendMessage(c, sendChannelId, calledCozeBotId, common.ImgGeneratePrompt+request.Prompt) + sentMsg, userAuth, err := discord.SendMessage(c, sendChannelId, calledCozeBotId, common.ImgGeneratePrompt+request.Prompt, maxToken) if err != nil { c.JSON(http.StatusInternalServerError, model.OpenAIErrorResponse{ OpenAIError: model.OpenAIError{ @@ -651,7 +651,7 @@ func ImagesForOpenAI(c *gin.Context) { } -func getSendChannelIdAndCozeBotId(c *gin.Context, channelId *string, model string, isOpenAIAPI bool) (sendChannelId string, calledCozeBotId string, isNewChannel bool, err error) { +func getSendChannelIdAndCozeBotId(c *gin.Context, channelId *string, model string, isOpenAIAPI bool) (sendChannelId string, calledCozeBotId string, maxToken string, isNewChannel bool, err error) { secret := "" if isOpenAIAPI { if secret = c.Request.Header.Get("Authorization"); secret != "" { @@ -669,50 +669,50 @@ func getSendChannelIdAndCozeBotId(c *gin.Context, channelId *string, model strin // 有值则随机一个 botConfig, err := common.RandomElement(botConfigs) if err != nil { - return "", "", false, err + return "", "", "", false, err } if channelId != nil && *channelId != "" { - return *channelId, botConfig.CozeBotId, false, nil + return *channelId, botConfig.CozeBotId, discord.MessageMaxToken, false, nil } if discord.DefaultChannelEnable == "1" { - return botConfig.ChannelId, botConfig.CozeBotId, false, nil + return botConfig.ChannelId, botConfig.CozeBotId, botConfig.MessageMaxToken, false, nil } else { var sendChannelId string sendChannelId, err := discord.CreateChannelWithRetry(c, discord.GuildId, fmt.Sprintf("cdp-chat-%s", c.Request.Context().Value(common.RequestIdKey)), 0) if err != nil { common.LogError(c, err.Error()) - return "", "", false, err + return "", "", "", false, err } - return sendChannelId, botConfig.CozeBotId, true, nil + return sendChannelId, botConfig.CozeBotId, botConfig.MessageMaxToken, true, nil } } // 没有值抛出异常 - return "", "", false, &myerr.ModelNotFoundError{ + return "", "", "", false, &myerr.ModelNotFoundError{ ErrCode: 500, Message: fmt.Sprintf("[proxy-secret:%s]+[model:%s]未匹配到有效bot", secret, model), } } else { if discord.BotConfigExist || discord.CozeBotId == "" { - return "", "", false, myerr.ErrNoBotId + return "", "", "", false, myerr.ErrNoBotId } if channelId != nil && *channelId != "" { - return *channelId, discord.CozeBotId, false, nil + return *channelId, discord.CozeBotId, discord.MessageMaxToken, false, nil } if discord.DefaultChannelEnable == "1" { - return discord.ChannelId, discord.CozeBotId, false, nil + return discord.ChannelId, discord.CozeBotId, discord.MessageMaxToken, false, nil } else { sendChannelId, err := discord.CreateChannelWithRetry(c, discord.GuildId, fmt.Sprintf("cdp-chat-%s", c.Request.Context().Value(common.RequestIdKey)), 0) if err != nil { //common.LogError(c, myerr.Error()) - return "", "", false, err + return "", "", "", false, err } - return sendChannelId, discord.CozeBotId, true, nil + return sendChannelId, discord.CozeBotId, discord.MessageMaxToken, true, nil } } } diff --git a/discord/discord.go b/discord/discord.go index 17bf37bd..098ca7d4 100644 --- a/discord/discord.go +++ b/discord/discord.go @@ -34,6 +34,7 @@ var CozeBotId = os.Getenv("COZE_BOT_ID") var GuildId = os.Getenv("GUILD_ID") var ChannelId = os.Getenv("CHANNEL_ID") var DefaultChannelEnable = os.Getenv("DEFAULT_CHANNEL_ENABLE") +var MessageMaxToken = os.Getenv("MESSAGE_MAX_TOKEN") var ProxyUrl = os.Getenv("PROXY_URL") var ChannelAutoDelTime = os.Getenv("CHANNEL_AUTO_DEL_TIME") var CozeBotStayActiveEnable = os.Getenv("COZE_BOT_STAY_ACTIVE_ENABLE") @@ -241,6 +242,15 @@ func checkEnvVariable() { } } + if MessageMaxToken == "" { + MessageMaxToken = strconv.Itoa(128 * 1000) + } else { + _, err := strconv.Atoi(MessageMaxToken) + if err != nil { + common.FatalLog("环境变量 MESSAGE_MAX_TOKEN 设置有误") + } + } + if telegram.NotifyTelegramBotToken != "" { err := telegram.InitTelegramBot() if err != nil { @@ -286,14 +296,20 @@ func loadBotConfig() { common.FatalLog("Error parsing JSON:", err) } - // 校验默认频道 - if DefaultChannelEnable == "1" { - for _, botConfig := range BotConfigList { - if botConfig.ChannelId == "" { - common.FatalLog("默认频道开关开启时,必须为每个Coze-Bot配置ChannelId") + for _, botConfig := range BotConfigList { + // 校验默认频道 + if DefaultChannelEnable == "1" && botConfig.ChannelId == "" { + common.FatalLog("默认频道开关开启时,必须为每个Coze-Bot配置ChannelId") + } + // 校验MaxToken + if botConfig.MessageMaxToken != "" { + _, err := strconv.Atoi(botConfig.MessageMaxToken) + if err != nil { + common.FatalLog(fmt.Sprintf("messageMaxToken 必须为数字!")) } } } + BotConfigExist = true common.SysLog(fmt.Sprintf("载入配置文件成功 BotConfigs: %+v", BotConfigList)) } @@ -468,7 +484,7 @@ func messageUpdate(s *discordgo.Session, m *discordgo.MessageUpdate) { return } -func SendMessage(c *gin.Context, channelID, cozeBotId, message string) (*discordgo.Message, string, error) { +func SendMessage(c *gin.Context, channelID, cozeBotId, message, maxToken string) (*discordgo.Message, string, error) { var ctx context.Context if c == nil { ctx = context.Background() @@ -490,7 +506,13 @@ func SendMessage(c *gin.Context, channelID, cozeBotId, message string) (*discord content = strings.Replace(content, `\u003e`, ">", -1) tokens := common.CountTokens(content) - if tokens > 128*1000 { + maxTokenInt, err := strconv.Atoi(maxToken) + if err != nil { + common.LogError(ctx, fmt.Sprintf("error sending message: %s", err)) + return &discordgo.Message{}, "", fmt.Errorf("error sending message") + } + + if tokens > maxTokenInt { common.LogError(ctx, fmt.Sprintf("prompt已超过限制,请分段发送 [%v] %s", tokens, content)) return nil, "", fmt.Errorf("prompt已超过限制,请分段发送 [%v]", tokens) } @@ -525,7 +547,7 @@ func SendMessage(c *gin.Context, channelID, cozeBotId, message string) (*discord if errors.As(err, &myErr) { // 无效则将此 auth 移除 UserAuthorizations = common.FilterSlice(UserAuthorizations, userAuth) - return SendMessage(c, channelID, cozeBotId, message) + return SendMessage(c, channelID, cozeBotId, message, maxToken) } common.LogError(ctx, fmt.Sprintf("error sending message: %s", err)) return nil, "", fmt.Errorf("error sending message") @@ -641,7 +663,7 @@ func stayActiveMessageTask() { common.SysError(fmt.Sprintf("ChannelId{%s} BotId{%s} 活跃机器人任务消息发送异常!雪花Id生成失败!", sendChannelId, config.CozeBotId)) continue } - _, _, err = SendMessage(nil, sendChannelId, config.CozeBotId, fmt.Sprintf("【%v】 %s", nextID, "CDP Scheduled Task Job Send Msg Success!")) + _, _, err = SendMessage(nil, sendChannelId, config.CozeBotId, fmt.Sprintf("【%v】 %s", nextID, "CDP Scheduled Task Job Send Msg Success!"), config.MessageMaxToken) if err != nil { common.SysError(fmt.Sprintf("ChannelId{%s} BotId{%s} 活跃机器人任务消息发送异常!", sendChannelId, config.CozeBotId)) } else { diff --git a/model/bot.go b/model/bot.go index 836a3081..78b806d2 100644 --- a/model/bot.go +++ b/model/bot.go @@ -1,10 +1,11 @@ package model type BotConfig struct { - ProxySecret string `json:"proxySecret"` - CozeBotId string `json:"cozeBotId"` - Model []string `json:"model"` - ChannelId string `json:"channelId"` + ProxySecret string `json:"proxySecret"` + CozeBotId string `json:"cozeBotId"` + Model []string `json:"model"` + ChannelId string `json:"channelId"` + MessageMaxToken string `json:"messageMaxToken"` } // FilterUniqueBotChannel 给定BotConfig切片,筛选出具有不同CozeBotId+ChannelId组合的元素