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

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

constant/context_key.go

@@ -18,8 +18,10 @@ const (
 	ContextKeyTokenSpecificChannelId ContextKey = "specific_channel_id"
 	ContextKeyTokenModelLimitEnabled ContextKey = "token_model_limit_enabled"
 	ContextKeyTokenModelLimit        ContextKey = "token_model_limit"
+	ContextKeyTokenCrossGroupRetry   ContextKey = "token_cross_group_retry"
 
 	/* channel related keys */
+	ContextKeyAutoGroupIndex           ContextKey = "auto_group_index"
 	ContextKeyChannelId                ContextKey = "channel_id"
 	ContextKeyChannelName              ContextKey = "channel_name"
 	ContextKeyChannelCreateTime        ContextKey = "channel_create_time"

controller/token.go

@@ -248,6 +248,7 @@ func UpdateToken(c *gin.Context) {
 		cleanToken.ModelLimits = token.ModelLimits
 		cleanToken.AllowIps = token.AllowIps
 		cleanToken.Group = token.Group
+		cleanToken.CrossGroupRetry = token.CrossGroupRetry
 	}
 	err = cleanToken.Update()
 	if err != nil {

middleware/auth.go

@@ -308,6 +308,7 @@ func SetupContextForToken(c *gin.Context, token *model.Token, parts ...string) e
 		c.Set("token_model_limit_enabled", false)
 	}
 	c.Set("token_group", token.Group)
+	c.Set("token_cross_group_retry", token.CrossGroupRetry)
 	if len(parts) > 1 {
 		if model.IsAdmin(token.UserId) {
 			c.Set("specific_channel_id", parts[1])

model/token.go

@@ -27,6 +27,7 @@ type Token struct {
 	AllowIps           *string        `json:"allow_ips" gorm:"default:''"`
 	UsedQuota          int            `json:"used_quota" gorm:"default:0"` // used quota
 	Group              string         `json:"group" gorm:"default:''"`
+	CrossGroupRetry    bool           `json:"cross_group_retry" gorm:"default:false"` // 跨分组重试，仅auto分组有效
 	DeletedAt          gorm.DeletedAt `gorm:"index"`
 }
 
@@ -185,7 +186,7 @@ func (token *Token) Update() (err error) {
 		}
 	}()
 	err = DB.Model(token).Select("name", "status", "expired_time", "remain_quota", "unlimited_quota",
-		"model_limits_enabled", "model_limits", "allow_ips", "group").Updates(token).Error
+		"model_limits_enabled", "model_limits", "allow_ips", "group", "cross_group_retry").Updates(token).Error
 	return err
 }
 

relay/channel/openai/helper.go

@@ -172,7 +172,7 @@ func handleLastResponse(lastStreamData string, responseId *string, createAt *int
 	shouldSendLastResp *bool) error {
 
 	var lastStreamResponse dto.ChatCompletionsStreamResponse
-	if err := json.Unmarshal(common.StringToByteSlice(lastStreamData), &lastStreamResponse); err != nil {
+	if err := common.Unmarshal(common.StringToByteSlice(lastStreamData), &lastStreamResponse); err != nil {
 		return err
 	}
 

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)

web/src/components/table/tokens/TokensColumnDefs.jsx

@@ -88,7 +88,7 @@ const renderStatus = (text, record, t) => {
 };
 
 // Render group column
-const renderGroupColumn = (text, t) => {
+const renderGroupColumn = (text, record, t) => {
   if (text === 'auto') {
     return (
       <Tooltip
@@ -98,8 +98,8 @@ const renderGroupColumn = (text, t) => {
         position='top'
       >
         <Tag color='white' shape='circle'>
-          {' '}
-          {t('智能熔断')}{' '}
+          {t('智能熔断')}
+          {record && record.cross_group_retry ? `(${t('跨分组')})` : ''}
         </Tag>
       </Tooltip>
     );
@@ -455,7 +455,7 @@ export const getTokensColumns = ({
       title: t('分组'),
       dataIndex: 'group',
       key: 'group',
-      render: (text) => renderGroupColumn(text, t),
+      render: (text, record) => renderGroupColumn(text, record, t),
     },
     {
       title: t('密钥'),

web/src/components/table/tokens/modals/EditTokenModal.jsx

@@ -73,6 +73,7 @@ const EditTokenModal = (props) => {
     model_limits: [],
     allow_ips: '',
     group: '',
+    cross_group_retry: false,
     tokenCount: 1,
   });
 
@@ -377,6 +378,16 @@ const EditTokenModal = (props) => {
                       />
                     )}
                   </Col>
+                  <Col span={24} style={{ display: values.group === 'auto' ? 'block' : 'none' }}>
+                    <Form.Switch
+                      field='cross_group_retry'
+                      label={t('跨分组重试')}
+                      size='default'
+                      extraText={t(
+                        '开启后，当前分组渠道失败时会按顺序尝试下一个分组的渠道',
+                      )}
+                    />
+                  </Col>
                   <Col xs={24} sm={24} md={24} lg={10} xl={10}>
                     <Form.DatePicker
                       field='expired_time'
@@ -499,7 +510,7 @@ const EditTokenModal = (props) => {
                     <Form.Switch
                       field='unlimited_quota'
                       label={t('无限额度')}
-                      size='large'
+                      size='default'
                       extraText={t(
                         '令牌的额度仅用于限制令牌本身的最大额度使用量，实际的使用受到账户的剩余额度限制',
                       )}

web/src/i18n/locales/en.json

@@ -2177,6 +2177,9 @@
     "默认区域，如: us-central1": "Default region, e.g.: us-central1",
     "默认折叠侧边栏": "Default collapse sidebar",
     "默认测试模型": "Default Test Model",
-    "默认补全倍率": "Default completion ratio"
+    "默认补全倍率": "Default completion ratio",
+    "跨分组重试": "Cross-group retry",
+    "跨分组": "Cross-group",
+    "开启后，当前分组渠道失败时会按顺序尝试下一个分组的渠道": "After enabling, when the current group channel fails, it will try the next group's channel in order"
   }
 }

web/src/i18n/locales/fr.json

@@ -2226,6 +2226,9 @@
     "默认助手消息": "Bonjour ! Comment puis-je vous aider aujourd'hui ?",
     "可选，用于复现结果": "Optionnel, pour des résultats reproductibles",
     "随机种子 (留空为随机)": "Graine aléatoire (laisser vide pour aléatoire)",
-    "默认补全倍率": "Taux de complétion par défaut"
+    "默认补全倍率": "Taux de complétion par défaut",
+    "跨分组重试": "Nouvelle tentative inter-groupes",
+    "跨分组": "Inter-groupes",
+    "开启后，当前分组渠道失败时会按顺序尝试下一个分组的渠道": "Après activation, lorsque le canal du groupe actuel échoue, il essaiera le canal du groupe suivant dans l'ordre"
   }
 }

web/src/i18n/locales/ja.json

@@ -2125,6 +2125,9 @@
     "默认用户消息": "こんにちは",
     "默认助手消息": "こんにちは！何かお手伝いできることはありますか？",
     "可选，用于复现结果": "オプション、結果の再現用",
-    "随机种子 (留空为随机)": "ランダムシード（空欄でランダム）"
+    "随机种子 (留空为随机)": "ランダムシード（空欄でランダム）",
+    "跨分组重试": "グループ間リトライ",
+    "跨分组": "グループ間",
+    "开启后，当前分组渠道失败时会按顺序尝试下一个分组的渠道": "有効にすると、現在のグループチャネルが失敗した場合、次のグループのチャネルを順番に試行します"
   }
 }

web/src/i18n/locales/ru.json

@@ -2236,6 +2236,9 @@
     "默认用户消息": "Здравствуйте",
     "默认助手消息": "Здравствуйте! Чем я могу вам помочь?",
     "可选，用于复现结果": "Необязательно, для воспроизводимых результатов",
-    "随机种子 (留空为随机)": "Случайное зерно (оставьте пустым для случайного)"
+    "随机种子 (留空为随机)": "Случайное зерно (оставьте пустым для случайного)",
+    "跨分组重试": "Повторная попытка между группами",
+    "跨分组": "Межгрупповой",
+    "开启后，当前分组渠道失败时会按顺序尝试下一个分组的渠道": "После включения, когда канал текущей группы не работает, он будет пытаться использовать канал следующей группы по порядку"
   }
 }

web/src/i18n/locales/vi.json

@@ -2736,6 +2736,9 @@
     "默认用户消息": "Xin chào",
     "默认助手消息": "Xin chào! Tôi có thể giúp gì cho bạn?",
     "可选，用于复现结果": "Tùy chọn, để tái tạo kết quả",
-    "随机种子 (留空为随机)": "Hạt giống ngẫu nhiên (để trống cho ngẫu nhiên)"
+    "随机种子 (留空为随机)": "Hạt giống ngẫu nhiên (để trống cho ngẫu nhiên)",
+    "跨分组重试": "Thử lại giữa các nhóm",
+    "跨分组": "Giữa các nhóm",
+    "开启后，当前分组渠道失败时会按顺序尝试下一个分组的渠道": "Sau khi bật, khi kênh nhóm hiện tại thất bại, nó sẽ thử kênh của nhóm tiếp theo theo thứ tự"
   }
 }

web/src/i18n/locales/zh.json

@@ -2203,6 +2203,9 @@
     "默认用户消息": "你好",
     "默认助手消息": "你好！有什么我可以帮助你的吗？",
     "可选，用于复现结果": "可选，用于复现结果",
-    "随机种子 (留空为随机)": "随机种子 (留空为随机)"
+    "随机种子 (留空为随机)": "随机种子 (留空为随机)",
+    "跨分组重试": "跨分组重试",
+    "跨分组": "跨分组",
+    "开启后，当前分组渠道失败时会按顺序尝试下一个分组的渠道": "开启后，当前分组渠道失败时会按顺序尝试下一个分组的渠道"
   }
 }