feat(token): add cross-group retry option for token processing
This commit is contained in:
@@ -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
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user