diff --git a/middleware/distributor.go b/middleware/distributor.go index cf5caa06d513..63c19df1feeb 100644 --- a/middleware/distributor.go +++ b/middleware/distributor.go @@ -83,22 +83,15 @@ func Distribute() func(c *gin.Context) { } var selectGroup string usingGroup := common.GetContextKeyString(c, constant.ContextKeyUsingGroup) - // check path is /pg/chat/completions - if strings.HasPrefix(c.Request.URL.Path, "/pg/chat/completions") { - playgroundRequest := &dto.PlayGroundRequest{} - err = common.UnmarshalBodyReusable(c, playgroundRequest) - if err != nil { - abortWithOpenAiMessage(c, http.StatusBadRequest, i18n.T(c, i18n.MsgDistributorInvalidPlayground, map[string]any{"Error": err.Error()})) + + // 支持从请求体中指定分组（标准 API 和 Playground 均支持） + if modelRequest.Group != "" { + if !service.GroupInUserUsableGroups(usingGroup, modelRequest.Group) && modelRequest.Group != usingGroup { + abortWithOpenAiMessage(c, http.StatusForbidden, i18n.T(c, i18n.MsgDistributorGroupAccessDenied)) return } - if playgroundRequest.Group != "" { - if !service.GroupInUserUsableGroups(usingGroup, playgroundRequest.Group) && playgroundRequest.Group != usingGroup { - abortWithOpenAiMessage(c, http.StatusForbidden, i18n.T(c, i18n.MsgDistributorGroupAccessDenied)) - return - } - usingGroup = playgroundRequest.Group - common.SetContextKey(c, constant.ContextKeyUsingGroup, usingGroup) - } + usingGroup = modelRequest.Group + common.SetContextKey(c, constant.ContextKeyUsingGroup, usingGroup) } if preferredChannelID, found := service.GetPreferredChannelByAffinity(c, modelRequest.Model, usingGroup); found { @@ -349,6 +342,7 @@ func getModelRequest(c *gin.Context) (*ModelRequest, bool, error) { return nil, false, err } modelRequest.Model = req.Model + modelRequest.Group = req.Group } if strings.HasPrefix(c.Request.URL.Path, "/v1/realtime") { //wss://api.openai.com/v1/realtime?model=gpt-4o-realtime-preview-2024-10-01 diff --git a/model/option.go b/model/option.go index 8e8587f271c8..45e82f02f214 100644 --- a/model/option.go +++ b/model/option.go @@ -145,6 +145,7 @@ func InitOptionMap() { common.OptionMap["CreateCacheRatio"] = ratio_setting.CreateCacheRatio2JSONString() common.OptionMap["GroupRatio"] = ratio_setting.GroupRatio2JSONString() common.OptionMap["GroupGroupRatio"] = ratio_setting.GroupGroupRatio2JSONString() + common.OptionMap["ModelGroupRatio"] = ratio_setting.ModelGroupRatio2JSONString() common.OptionMap["UserUsableGroups"] = setting.UserUsableGroups2JSONString() common.OptionMap["CompletionRatio"] = ratio_setting.CompletionRatio2JSONString() common.OptionMap["ImageRatio"] = ratio_setting.ImageRatio2JSONString() @@ -530,6 +531,8 @@ func updateOptionMap(key string, value string) (err error) { err = ratio_setting.UpdateGroupRatioByJSONString(value) case "GroupGroupRatio": err = ratio_setting.UpdateGroupGroupRatioByJSONString(value) + case "ModelGroupRatio": + err = ratio_setting.UpdateModelGroupRatioByJSONString(value) case "UserUsableGroups": err = setting.UpdateUserUsableGroupsByJSONString(value) case "CompletionRatio": diff --git a/relay/helper/price.go b/relay/helper/price.go index 2e8ebb2d2fc1..11d7c64f2523 100644 --- a/relay/helper/price.go +++ b/relay/helper/price.go @@ -66,6 +66,12 @@ func HandleGroupRatio(ctx *gin.Context, relayInfo *relaycommon.RelayInfo) types. groupRatioInfo.GroupRatio = ratio_setting.GetGroupRatio(relayInfo.UsingGroup) } + // apply model-specific group ratio overlay + modelName := relayInfo.OriginModelName + if modelName != "" { + groupRatioInfo.GroupRatio = ratio_setting.GetModelGroupRatio(modelName, relayInfo.UsingGroup, groupRatioInfo.GroupRatio) + } + return groupRatioInfo } diff --git a/service/channel_select.go b/service/channel_select.go index 24c4e252bfb3..a0e127939c89 100644 --- a/service/channel_select.go +++ b/service/channel_select.go @@ -2,6 +2,7 @@ package service import ( "errors" + "strings" "github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/constant" @@ -49,54 +50,68 @@ func (p *RetryParam) ResetRetryNextTry() { // CacheGetRandomSatisfiedChannel tries to get a random channel that satisfies the requirements. // 尝试获取一个满足要求的随机渠道。 // -// For "auto" tokenGroup with cross-group Retry enabled: -// 对于启用了跨分组重试的 "auto" tokenGroup： +// Supports three modes for TokenGroup: +// 支持三种 TokenGroup 模式： // -// - Each group will exhaust all its priorities before moving to the next group. -// 每个分组会用完所有优先级后才会切换到下一个分组。 -// -// - Uses ContextKeyAutoGroupIndex to track current group index. -// 使用 ContextKeyAutoGroupIndex 跟踪当前分组索引。 -// -// - Uses ContextKeyAutoGroupRetryIndex to track the global Retry count when current group started. -// 使用 ContextKeyAutoGroupRetryIndex 跟踪当前分组开始时的全局重试次数。 +// 1. Single group (e.g. "vip") — uses that group only +// 单一分组 — 仅使用该分组 // -// - priorityRetry = Retry - startRetryIndex, represents the priority level within current group. -// priorityRetry = Retry - startRetryIndex，表示当前分组内的优先级级别。 +// 2. "auto" — uses globally configured auto_groups +// "auto" — 使用全局配置的自动分组 // -// - When GetRandomSatisfiedChannel returns nil (priorities exhausted), moves to next group. -// 当 GetRandomSatisfiedChannel 返回 nil（优先级用完）时，切换到下一个分组。 +// 3. Comma-separated list (e.g. "cheap,premium,fallback") — custom fallback order +// 逗号分隔列表 — 自定义 fallback 顺序，按列表顺序依次尝试 // -// Example flow (2 groups, each with 2 priorities, RetryTimes=3): -// 示例流程（2个分组，每个有2个优先级，RetryTimes=3）： +// For auto and custom-fallback modes with cross-group Retry: +// 对于 auto 和自定义 fallback 模式的跨分组重试： // -// Retry=0: GroupA, priority0 (startRetryIndex=0, priorityRetry=0) -// 分组A, 优先级0 +// - Each group will exhaust all its priorities before moving to the next group. +// 每个分组会用完所有优先级后才会切换到下一个分组。 // -// Retry=1: GroupA, priority1 (startRetryIndex=0, priorityRetry=1) -// 分组A, 优先级1 +// - Uses ContextKeyAutoGroupIndex to track current group index. +// 使用 ContextKeyAutoGroupIndex 跟踪当前分组索引。 // -// Retry=2: GroupA exhausted → GroupB, priority0 (startRetryIndex=2, priorityRetry=0) -// 分组A用完 → 分组B, 优先级0 +// Example flow (custom fallback "cheap,premium", each with 2 priorities, RetryTimes=3): +// 示例流程（自定义 fallback "cheap,premium"，每个有 2 个优先级，RetryTimes=3）： // -// Retry=3: GroupB, priority1 (startRetryIndex=2, priorityRetry=1) -// 分组B, 优先级1 +// Retry=0: cheap, priority0 +// Retry=1: cheap, priority1 +// Retry=2: cheap exhausted → premium, priority0 +// Retry=3: premium, priority1 func CacheGetRandomSatisfiedChannel(param *RetryParam) (*model.Channel, string, error) { var channel *model.Channel var err error selectGroup := param.TokenGroup userGroup := common.GetContextKeyString(param.Ctx, constant.ContextKeyUserGroup) - if param.TokenGroup == "auto" { + // 判断是否为多分组模式（auto 或逗号分隔的自定义 fallback 列表） + var groups []string + customFallback := false + if strings.Contains(param.TokenGroup, ",") { + // 自定义 fallback：逗号分隔的分组列表，如 "cheap,premium,fallback" + for _, g := range strings.Split(param.TokenGroup, ",") { + g = strings.TrimSpace(g) + if g != "" { + groups = append(groups, g) + } + } + customFallback = true + } else if param.TokenGroup == "auto" { if len(setting.GetAutoGroups()) == 0 { return nil, selectGroup, errors.New("auto groups is not enabled") } - autoGroups := GetUserAutoGroup(userGroup) + groups = GetUserAutoGroup(userGroup) + } - // startGroupIndex: the group index to start searching from - // startGroupIndex: 开始搜索的分组索引 + // 记录最后一个非 nil 错误，当所有分组都失败时返回 + var lastErr error + if len(groups) > 0 { startGroupIndex := 0 crossGroupRetry := common.GetContextKeyBool(param.Ctx, constant.ContextKeyTokenCrossGroupRetry) + // 自定义 fallback 模式默认启用跨分组重试 + if customFallback { + crossGroupRetry = true + } if lastGroupIndex, exists := common.GetContextKey(param.Ctx, constant.ContextKeyAutoGroupIndex); exists { if idx, ok := lastGroupIndex.(int); ok { @@ -104,51 +119,35 @@ func CacheGetRandomSatisfiedChannel(param *RetryParam) (*model.Channel, string, } } - for i := startGroupIndex; i < len(autoGroups); i++ { - autoGroup := autoGroups[i] - // Calculate priorityRetry for current group - // 计算当前分组的 priorityRetry + for i := startGroupIndex; i < len(groups); i++ { + group := groups[i] priorityRetry := param.GetRetry() - // If moved to a new group, reset priorityRetry and update startRetryIndex - // 如果切换到新分组，重置 priorityRetry 并更新 startRetryIndex if i > startGroupIndex { priorityRetry = 0 } - logger.LogDebug(param.Ctx, "Auto selecting group: %s, priorityRetry: %d", autoGroup, priorityRetry) + logger.LogDebug(param.Ctx, "Auto selecting group: %s, priorityRetry: %d", group, priorityRetry) - channel, _ = model.GetRandomSatisfiedChannel(autoGroup, param.ModelName, priorityRetry, param.RequestPath) + channel, err = model.GetRandomSatisfiedChannel(group, param.ModelName, priorityRetry, param.RequestPath) + if err != nil { + lastErr = err + } if channel == nil { - // Current group has no available channel for this model, try next group - // 当前分组没有该模型的可用渠道，尝试下一个分组 - logger.LogDebug(param.Ctx, "No available channel in group %s for model %s at priorityRetry %d, trying next group", autoGroup, param.ModelName, priorityRetry) - // 重置状态以尝试下一个分组 + logger.LogDebug(param.Ctx, "No available channel in group %s for model %s at priorityRetry %d, trying next group", group, param.ModelName, priorityRetry) common.SetContextKey(param.Ctx, constant.ContextKeyAutoGroupIndex, i+1) common.SetContextKey(param.Ctx, constant.ContextKeyAutoGroupRetryIndex, 0) - // Reset retry counter so outer loop can continue for next group - // 重置重试计数器，以便外层循环可以为下一个分组继续 param.SetRetry(0) continue } - common.SetContextKey(param.Ctx, constant.ContextKeyAutoGroup, autoGroup) - selectGroup = autoGroup - logger.LogDebug(param.Ctx, "Auto selected group: %s", autoGroup) + common.SetContextKey(param.Ctx, constant.ContextKeyAutoGroup, group) + selectGroup = group + logger.LogDebug(param.Ctx, "Auto selected group: %s", group) - // Prepare state for next retry - // 为下一次重试准备状态 if crossGroupRetry && priorityRetry >= common.RetryTimes { - // Current group has exhausted all retries, prepare to switch to next group - // This request still uses current group, but next retry will use next group - // 当前分组已用完所有重试次数，准备切换到下一个分组 - // 本次请求仍使用当前分组，但下次重试将使用下一个分组 - logger.LogDebug(param.Ctx, "Current group %s retries exhausted (priorityRetry=%d >= RetryTimes=%d), preparing switch to next group for next retry", autoGroup, priorityRetry, common.RetryTimes) + logger.LogDebug(param.Ctx, "Current group %s retries exhausted (priorityRetry=%d >= RetryTimes=%d), preparing switch to next group for next retry", group, priorityRetry, common.RetryTimes) common.SetContextKey(param.Ctx, constant.ContextKeyAutoGroupIndex, i+1) - // Reset retry counter so outer loop can continue for next group - // 重置重试计数器，以便外层循环可以为下一个分组继续 param.SetRetry(0) param.ResetRetryNextTry() } else { - // Stay in current group, save current state - // 保持在当前分组，保存当前状态 common.SetContextKey(param.Ctx, constant.ContextKeyAutoGroupIndex, i) } break @@ -159,5 +158,9 @@ func CacheGetRandomSatisfiedChannel(param *RetryParam) (*model.Channel, string, return nil, param.TokenGroup, err } } + // 所有分组都未找到可用渠道，返回记录的错误 + if channel == nil && lastErr != nil { + return nil, selectGroup, lastErr + } return channel, selectGroup, nil } diff --git a/service/group.go b/service/group.go index a73642c3eb1a..3ba253b8290d 100644 --- a/service/group.go +++ b/service/group.go @@ -63,3 +63,10 @@ func GetUserGroupRatio(userGroup, group string) float64 { } return ratio_setting.GetGroupRatio(group) } + +// GetModelGroupRatio 获取模型在指定分组中的倍率（模型×分组交叉倍率） +// 如果模型在该分组有自定义倍率则返回，否则回退到全局分组倍率 +func GetModelGroupRatio(modelName, group string) float64 { + globalGroupRatio := ratio_setting.GetGroupRatio(group) + return ratio_setting.GetModelGroupRatio(modelName, group, globalGroupRatio) +} diff --git a/setting/ratio_setting/group_ratio.go b/setting/ratio_setting/group_ratio.go index 7d16d9283932..cfc655ba5251 100644 --- a/setting/ratio_setting/group_ratio.go +++ b/setting/ratio_setting/group_ratio.go @@ -27,10 +27,18 @@ var groupGroupRatioMap = types.NewRWMap[string, map[string]float64]() var defaultGroupSpecialUsableGroup = map[string]map[string]string{} +// defaultModelGroupRatio 模型×分组交叉倍率 +// 格式: modelName -> groupName -> ratio +// 当模型在特定分组中有自定义倍率时，覆盖全局 group_ratio +var defaultModelGroupRatio = map[string]map[string]float64{} + +var modelGroupRatioMap = types.NewRWMap[string, map[string]float64]() + type GroupRatioSetting struct { GroupRatio *types.RWMap[string, float64] `json:"group_ratio"` GroupGroupRatio *types.RWMap[string, map[string]float64] `json:"group_group_ratio"` GroupSpecialUsableGroup *types.RWMap[string, map[string]string] `json:"group_special_usable_group"` + ModelGroupRatio *types.RWMap[string, map[string]float64] `json:"model_group_ratio"` } var groupRatioSetting GroupRatioSetting @@ -41,11 +49,13 @@ func init() { groupRatioMap.AddAll(defaultGroupRatio) groupGroupRatioMap.AddAll(defaultGroupGroupRatio) + modelGroupRatioMap.AddAll(defaultModelGroupRatio) groupRatioSetting = GroupRatioSetting{ GroupSpecialUsableGroup: groupSpecialUsableGroup, GroupRatio: groupRatioMap, GroupGroupRatio: groupGroupRatioMap, + ModelGroupRatio: modelGroupRatioMap, } config.GlobalConfig.Register("group_ratio_setting", &groupRatioSetting) @@ -118,3 +128,29 @@ func CheckGroupRatio(jsonStr string) error { } return nil } + +// GetModelGroupRatio 获取模型在指定分组中的倍率 +// 如果模型在该分组有自定义倍率则返回，否则返回传入的 fallbackGroupRatio +func GetModelGroupRatio(modelName, group string, fallbackGroupRatio float64) float64 { + modelRatios, ok := modelGroupRatioMap.Get(modelName) + if !ok { + return fallbackGroupRatio + } + ratio, ok := modelRatios[group] + if !ok { + return fallbackGroupRatio + } + return ratio +} + +func GetModelGroupRatioCopy() map[string]map[string]float64 { + return modelGroupRatioMap.ReadAll() +} + +func ModelGroupRatio2JSONString() string { + return modelGroupRatioMap.MarshalJSONString() +} + +func UpdateModelGroupRatioByJSONString(jsonStr string) error { + return types.LoadFromJsonString(modelGroupRatioMap, jsonStr) +}