Skip to content

Commit 2db58ce

Browse files
committed
feat: MESSAGE_MAX_TOKEN
1 parent 887e67c commit 2db58ce

File tree

4 files changed

+53
-30
lines changed

4 files changed

+53
-30
lines changed

common/constants.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
package common
22

3-
var Version = "v4.5.10" // this hard coding will be replaced automatically when building, no need to manually change
3+
var Version = "v4.6.0" // this hard coding will be replaced automatically when building, no need to manually change
44

55
const (
66
RequestIdKey = "X-Request-Id"

controller/chat.go

Lines changed: 16 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -166,7 +166,7 @@ func ChatForOpenAI(c *gin.Context) {
166166
return
167167
}
168168

169-
sendChannelId, calledCozeBotId, isNewChannel, err := getSendChannelIdAndCozeBotId(c, request.ChannelId, request.Model, true)
169+
sendChannelId, calledCozeBotId, maxToken, isNewChannel, err := getSendChannelIdAndCozeBotId(c, request.ChannelId, request.Model, true)
170170

171171
if err != nil {
172172
response := model.OpenAIErrorResponse{
@@ -275,7 +275,7 @@ loop:
275275
content = string(jsonData)
276276
}
277277

278-
sentMsg, userAuth, err := discord.SendMessage(c, sendChannelId, calledCozeBotId, content)
278+
sentMsg, userAuth, err := discord.SendMessage(c, sendChannelId, calledCozeBotId, content, maxToken)
279279
if err != nil {
280280
c.JSON(http.StatusInternalServerError, model.OpenAIErrorResponse{
281281
OpenAIError: model.OpenAIError{
@@ -517,7 +517,7 @@ func ImagesForOpenAI(c *gin.Context) {
517517
return
518518
}
519519

520-
sendChannelId, calledCozeBotId, isNewChannel, err := getSendChannelIdAndCozeBotId(c, request.ChannelId, request.Model, true)
520+
sendChannelId, calledCozeBotId, maxToken, isNewChannel, err := getSendChannelIdAndCozeBotId(c, request.ChannelId, request.Model, true)
521521
if err != nil {
522522
common.LogError(c.Request.Context(), err.Error())
523523
c.JSON(http.StatusInternalServerError, model.OpenAIErrorResponse{
@@ -547,7 +547,7 @@ func ImagesForOpenAI(c *gin.Context) {
547547
}()
548548
}
549549

550-
sentMsg, userAuth, err := discord.SendMessage(c, sendChannelId, calledCozeBotId, common.ImgGeneratePrompt+request.Prompt)
550+
sentMsg, userAuth, err := discord.SendMessage(c, sendChannelId, calledCozeBotId, common.ImgGeneratePrompt+request.Prompt, maxToken)
551551
if err != nil {
552552
c.JSON(http.StatusInternalServerError, model.OpenAIErrorResponse{
553553
OpenAIError: model.OpenAIError{
@@ -651,7 +651,7 @@ func ImagesForOpenAI(c *gin.Context) {
651651

652652
}
653653

654-
func getSendChannelIdAndCozeBotId(c *gin.Context, channelId *string, model string, isOpenAIAPI bool) (sendChannelId string, calledCozeBotId string, isNewChannel bool, err error) {
654+
func getSendChannelIdAndCozeBotId(c *gin.Context, channelId *string, model string, isOpenAIAPI bool) (sendChannelId string, calledCozeBotId string, maxToken string, isNewChannel bool, err error) {
655655
secret := ""
656656
if isOpenAIAPI {
657657
if secret = c.Request.Header.Get("Authorization"); secret != "" {
@@ -669,50 +669,50 @@ func getSendChannelIdAndCozeBotId(c *gin.Context, channelId *string, model strin
669669
// 有值则随机一个
670670
botConfig, err := common.RandomElement(botConfigs)
671671
if err != nil {
672-
return "", "", false, err
672+
return "", "", "", false, err
673673
}
674674

675675
if channelId != nil && *channelId != "" {
676-
return *channelId, botConfig.CozeBotId, false, nil
676+
return *channelId, botConfig.CozeBotId, discord.MessageMaxToken, false, nil
677677
}
678678

679679
if discord.DefaultChannelEnable == "1" {
680-
return botConfig.ChannelId, botConfig.CozeBotId, false, nil
680+
return botConfig.ChannelId, botConfig.CozeBotId, botConfig.MessageMaxToken, false, nil
681681
} else {
682682
var sendChannelId string
683683
sendChannelId, err := discord.CreateChannelWithRetry(c, discord.GuildId, fmt.Sprintf("cdp-chat-%s", c.Request.Context().Value(common.RequestIdKey)), 0)
684684
if err != nil {
685685
common.LogError(c, err.Error())
686-
return "", "", false, err
686+
return "", "", "", false, err
687687
}
688-
return sendChannelId, botConfig.CozeBotId, true, nil
688+
return sendChannelId, botConfig.CozeBotId, botConfig.MessageMaxToken, true, nil
689689
}
690690

691691
}
692692
// 没有值抛出异常
693-
return "", "", false, &myerr.ModelNotFoundError{
693+
return "", "", "", false, &myerr.ModelNotFoundError{
694694
ErrCode: 500,
695695
Message: fmt.Sprintf("[proxy-secret:%s]+[model:%s]未匹配到有效bot", secret, model),
696696
}
697697
} else {
698698

699699
if discord.BotConfigExist || discord.CozeBotId == "" {
700-
return "", "", false, myerr.ErrNoBotId
700+
return "", "", "", false, myerr.ErrNoBotId
701701
}
702702

703703
if channelId != nil && *channelId != "" {
704-
return *channelId, discord.CozeBotId, false, nil
704+
return *channelId, discord.CozeBotId, discord.MessageMaxToken, false, nil
705705
}
706706

707707
if discord.DefaultChannelEnable == "1" {
708-
return discord.ChannelId, discord.CozeBotId, false, nil
708+
return discord.ChannelId, discord.CozeBotId, discord.MessageMaxToken, false, nil
709709
} else {
710710
sendChannelId, err := discord.CreateChannelWithRetry(c, discord.GuildId, fmt.Sprintf("cdp-chat-%s", c.Request.Context().Value(common.RequestIdKey)), 0)
711711
if err != nil {
712712
//common.LogError(c, myerr.Error())
713-
return "", "", false, err
713+
return "", "", "", false, err
714714
}
715-
return sendChannelId, discord.CozeBotId, true, nil
715+
return sendChannelId, discord.CozeBotId, discord.MessageMaxToken, true, nil
716716
}
717717
}
718718
}

discord/discord.go

Lines changed: 31 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@ var CozeBotId = os.Getenv("COZE_BOT_ID")
3434
var GuildId = os.Getenv("GUILD_ID")
3535
var ChannelId = os.Getenv("CHANNEL_ID")
3636
var DefaultChannelEnable = os.Getenv("DEFAULT_CHANNEL_ENABLE")
37+
var MessageMaxToken = os.Getenv("MESSAGE_MAX_TOKEN")
3738
var ProxyUrl = os.Getenv("PROXY_URL")
3839
var ChannelAutoDelTime = os.Getenv("CHANNEL_AUTO_DEL_TIME")
3940
var CozeBotStayActiveEnable = os.Getenv("COZE_BOT_STAY_ACTIVE_ENABLE")
@@ -241,6 +242,15 @@ func checkEnvVariable() {
241242
}
242243
}
243244

245+
if MessageMaxToken == "" {
246+
MessageMaxToken = strconv.Itoa(128 * 1000)
247+
} else {
248+
_, err := strconv.Atoi(MessageMaxToken)
249+
if err != nil {
250+
common.FatalLog("环境变量 MESSAGE_MAX_TOKEN 设置有误")
251+
}
252+
}
253+
244254
if telegram.NotifyTelegramBotToken != "" {
245255
err := telegram.InitTelegramBot()
246256
if err != nil {
@@ -286,14 +296,20 @@ func loadBotConfig() {
286296
common.FatalLog("Error parsing JSON:", err)
287297
}
288298

289-
// 校验默认频道
290-
if DefaultChannelEnable == "1" {
291-
for _, botConfig := range BotConfigList {
292-
if botConfig.ChannelId == "" {
293-
common.FatalLog("默认频道开关开启时,必须为每个Coze-Bot配置ChannelId")
299+
for _, botConfig := range BotConfigList {
300+
// 校验默认频道
301+
if DefaultChannelEnable == "1" && botConfig.ChannelId == "" {
302+
common.FatalLog("默认频道开关开启时,必须为每个Coze-Bot配置ChannelId")
303+
}
304+
// 校验MaxToken
305+
if botConfig.MessageMaxToken != "" {
306+
_, err := strconv.Atoi(botConfig.MessageMaxToken)
307+
if err != nil {
308+
common.FatalLog(fmt.Sprintf("messageMaxToken 必须为数字!"))
294309
}
295310
}
296311
}
312+
297313
BotConfigExist = true
298314
common.SysLog(fmt.Sprintf("载入配置文件成功 BotConfigs: %+v", BotConfigList))
299315
}
@@ -468,7 +484,7 @@ func messageUpdate(s *discordgo.Session, m *discordgo.MessageUpdate) {
468484
return
469485
}
470486

471-
func SendMessage(c *gin.Context, channelID, cozeBotId, message string) (*discordgo.Message, string, error) {
487+
func SendMessage(c *gin.Context, channelID, cozeBotId, message, maxToken string) (*discordgo.Message, string, error) {
472488
var ctx context.Context
473489
if c == nil {
474490
ctx = context.Background()
@@ -490,7 +506,13 @@ func SendMessage(c *gin.Context, channelID, cozeBotId, message string) (*discord
490506
content = strings.Replace(content, `\u003e`, ">", -1)
491507

492508
tokens := common.CountTokens(content)
493-
if tokens > 128*1000 {
509+
maxTokenInt, err := strconv.Atoi(maxToken)
510+
if err != nil {
511+
common.LogError(ctx, fmt.Sprintf("error sending message: %s", err))
512+
return &discordgo.Message{}, "", fmt.Errorf("error sending message")
513+
}
514+
515+
if tokens > maxTokenInt {
494516
common.LogError(ctx, fmt.Sprintf("prompt已超过限制,请分段发送 [%v] %s", tokens, content))
495517
return nil, "", fmt.Errorf("prompt已超过限制,请分段发送 [%v]", tokens)
496518
}
@@ -525,7 +547,7 @@ func SendMessage(c *gin.Context, channelID, cozeBotId, message string) (*discord
525547
if errors.As(err, &myErr) {
526548
// 无效则将此 auth 移除
527549
UserAuthorizations = common.FilterSlice(UserAuthorizations, userAuth)
528-
return SendMessage(c, channelID, cozeBotId, message)
550+
return SendMessage(c, channelID, cozeBotId, message, maxToken)
529551
}
530552
common.LogError(ctx, fmt.Sprintf("error sending message: %s", err))
531553
return nil, "", fmt.Errorf("error sending message")
@@ -641,7 +663,7 @@ func stayActiveMessageTask() {
641663
common.SysError(fmt.Sprintf("ChannelId{%s} BotId{%s} 活跃机器人任务消息发送异常!雪花Id生成失败!", sendChannelId, config.CozeBotId))
642664
continue
643665
}
644-
_, _, err = SendMessage(nil, sendChannelId, config.CozeBotId, fmt.Sprintf("【%v】 %s", nextID, "CDP Scheduled Task Job Send Msg Success!"))
666+
_, _, err = SendMessage(nil, sendChannelId, config.CozeBotId, fmt.Sprintf("【%v】 %s", nextID, "CDP Scheduled Task Job Send Msg Success!"), config.MessageMaxToken)
645667
if err != nil {
646668
common.SysError(fmt.Sprintf("ChannelId{%s} BotId{%s} 活跃机器人任务消息发送异常!", sendChannelId, config.CozeBotId))
647669
} else {

model/bot.go

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,11 @@
11
package model
22

33
type BotConfig struct {
4-
ProxySecret string `json:"proxySecret"`
5-
CozeBotId string `json:"cozeBotId"`
6-
Model []string `json:"model"`
7-
ChannelId string `json:"channelId"`
4+
ProxySecret string `json:"proxySecret"`
5+
CozeBotId string `json:"cozeBotId"`
6+
Model []string `json:"model"`
7+
ChannelId string `json:"channelId"`
8+
MessageMaxToken string `json:"messageMaxToken"`
89
}
910

1011
// FilterUniqueBotChannel 给定BotConfig切片,筛选出具有不同CozeBotId+ChannelId组合的元素

0 commit comments

Comments
 (0)