Skip to content

Commit

Permalink
feat: MESSAGE_MAX_TOKEN
Browse files Browse the repository at this point in the history
  • Loading branch information
deanxv committed Aug 17, 2024
1 parent 887e67c commit 2db58ce
Show file tree
Hide file tree
Showing 4 changed files with 53 additions and 30 deletions.
2 changes: 1 addition & 1 deletion common/constants.go
Original file line number Diff line number Diff line change
@@ -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"
Expand Down
32 changes: 16 additions & 16 deletions controller/chat.go
Original file line number Diff line number Diff line change
Expand Up @@ -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{
Expand Down Expand Up @@ -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{
Expand Down Expand Up @@ -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{
Expand Down Expand Up @@ -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{
Expand Down Expand Up @@ -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 != "" {
Expand All @@ -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
}
}
}
Expand Down
40 changes: 31 additions & 9 deletions discord/discord.go
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -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))
}
Expand Down Expand Up @@ -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()
Expand All @@ -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)
}
Expand Down Expand Up @@ -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")
Expand Down Expand Up @@ -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 {
Expand Down
9 changes: 5 additions & 4 deletions model/bot.go
Original file line number Diff line number Diff line change
@@ -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组合的元素
Expand Down

0 comments on commit 2db58ce

Please sign in to comment.