github/new-api
1
0
Fork 0
mirror of https://github.com/QuantumNous/new-api.git synced 2026-04-19 05:57:28 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(token): add cross-group retry option for token processing

Browse Source
This commit is contained in:
CaIon
2025-12-12 17:59:21 +08:00
parent 0b9f6a58bc
commit 01b4039e96
8 changed files with 38 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

22
service/channel_select.go
View File

@@ -11,6 +11,7 @@ import (
"github.com/gin-gonic/gin"
)
// CacheGetRandomSatisfiedChannel tries to get a random channel that satisfies the requirements.
func CacheGetRandomSatisfiedChannel(c *gin.Context, group string, modelName string, retry int) (*model.Channel, string, error) {
var channel *model.Channel
var err error
@@ -20,15 +21,28 @@ func CacheGetRandomSatisfiedChannel(c *gin.Context, group string, modelName stri
if len(setting.GetAutoGroups()) == 0 {
return nil, selectGroup, errors.New("auto groups is not enabled")
}
for _, autoGroup := range GetUserAutoGroup(userGroup) {
logger.LogDebug(c, "Auto selecting group:", autoGroup)
channel, _ = model.GetRandomSatisfiedChannel(autoGroup, modelName, retry)
autoGroups := GetUserAutoGroup(userGroup)
// 如果 token 启用了跨分组重试，获取上次失败的 auto group 索引，从下一个开始尝试
startIndex := 0
crossGroupRetry := common.GetContextKeyBool(c, constant.ContextKeyTokenCrossGroupRetry)
if crossGroupRetry && retry > 0 {
logger.LogDebug(c, "Auto group retry cross group, retry: %d", retry)
if lastIndex, exists := c.Get(string(constant.ContextKeyAutoGroupIndex)); exists {
startIndex = lastIndex.(int) + 1
}
logger.LogDebug(c, "Auto group retry cross group, start index: %d", startIndex)
}
for i := startIndex; i < len(autoGroups); i++ {
autoGroup := autoGroups[i]
logger.LogDebug(c, "Auto selecting group: %s", autoGroup)
channel, _ = model.GetRandomSatisfiedChannel(autoGroup, modelName, 0)
if channel == nil {
continue
} else {
c.Set("auto_group", autoGroup)
c.Set(string(constant.ContextKeyAutoGroupIndex), i)
selectGroup = autoGroup
logger.LogDebug(c, "Auto selected group:", autoGroup)
logger.LogDebug(c, "Auto selected group: %s", autoGroup)
break
}
}

2
service/token_counter.go
View File

@@ -317,7 +317,7 @@ func EstimateRequestToken(c *gin.Context, meta *types.TokenCountMeta, info *rela
for i, file := range meta.Files {
switch file.FileType {
case types.FileTypeImage:
if common.IsOpenAITextModel(info.OriginModelName) {
if common.IsOpenAITextModel(model) {
token, err := getImageToken(file, model, info.IsStream)
if err != nil {
return 0, fmt.Errorf("error counting image token, media index[%d], original data[%s], err: %v", i, file.OriginData, err)