Merge pull request #2426 from QuantumNous/feat/auto-cross-group-retry

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

service/channel_select.go

@@ -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,30 @@ 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 := common.GetContextKey(c, constant.ContextKeyAutoGroupIndex); exists {
+				if idx, ok := lastIndex.(int); ok {
+					startIndex = idx + 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)
+				common.SetContextKey(c, constant.ContextKeyAutoGroupIndex, i)
 				selectGroup = autoGroup
-				logger.LogDebug(c, "Auto selected group:", autoGroup)
+				logger.LogDebug(c, "Auto selected group: %s", autoGroup)
 				break
 			}
 		}

service/token_counter.go

@@ -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)