From b49842d5432c255693f40ddba0bd4cbeb7aa3915 Mon Sep 17 00:00:00 2001 From: fengsilin Date: Sat, 16 May 2026 01:20:52 +0800 Subject: [PATCH] =?UTF-8?q?refactor:=20=E6=B8=85=E7=90=86=20Distribute=20?= =?UTF-8?q?=E4=B8=AD=E7=9A=84=E5=86=97=E4=BD=99=E6=9D=A1=E4=BB=B6=E5=88=A4?= =?UTF-8?q?=E6=96=AD=E5=92=8C=E7=BC=A9=E8=BF=9B?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 删除 tautological 的 `if channel == nil` 外层判断(channel 在该位置永远为 nil) - 修正删除指定渠道分支后遗留的多余缩进(减少一级 tab) - 更新通道亲和性检查注释为简洁版 --- middleware/distributor.go | 156 +++++++++++++++++++------------------- 1 file changed, 77 insertions(+), 79 deletions(-) diff --git a/middleware/distributor.go b/middleware/distributor.go index 3938379..5a9bbf5 100644 --- a/middleware/distributor.go +++ b/middleware/distributor.go @@ -34,102 +34,100 @@ func Distribute() func(c *gin.Context) { abortWithOpenAiMessage(c, http.StatusBadRequest, i18n.T(c, i18n.MsgDistributorInvalidRequest, map[string]any{"Error": err.Error()})) return } - - // Normal channel selection logic - // check token model mapping - modelLimitEnable := common.GetContextKeyBool(c, constant.ContextKeyTokenModelLimitEnabled) - if modelLimitEnable { - s, ok := common.GetContextKey(c, constant.ContextKeyTokenModelLimit) - if !ok { - // token model limit is empty, all models are not allowed - abortWithOpenAiMessage(c, http.StatusForbidden, i18n.T(c, i18n.MsgDistributorTokenNoModelAccess)) - return - } - var tokenModelLimit map[string]bool - tokenModelLimit, ok = s.(map[string]bool) - if !ok { - tokenModelLimit = map[string]bool{} - } - matchName := ratio_setting.FormatMatchingModelName(modelRequest.Model) // match gpts & thinking-* - if _, ok := tokenModelLimit[matchName]; !ok { - abortWithOpenAiMessage(c, http.StatusForbidden, i18n.T(c, i18n.MsgDistributorTokenModelForbidden, map[string]any{"Model": modelRequest.Model})) - return - } + + // Normal channel selection logic + // check token model mapping + modelLimitEnable := common.GetContextKeyBool(c, constant.ContextKeyTokenModelLimitEnabled) + if modelLimitEnable { + s, ok := common.GetContextKey(c, constant.ContextKeyTokenModelLimit) + if !ok { + // token model limit is empty, all models are not allowed + abortWithOpenAiMessage(c, http.StatusForbidden, i18n.T(c, i18n.MsgDistributorTokenNoModelAccess)) + return + } + var tokenModelLimit map[string]bool + tokenModelLimit, ok = s.(map[string]bool) + if !ok { + tokenModelLimit = map[string]bool{} } + matchName := ratio_setting.FormatMatchingModelName(modelRequest.Model) // match gpts & thinking-* + if _, ok := tokenModelLimit[matchName]; !ok { + abortWithOpenAiMessage(c, http.StatusForbidden, i18n.T(c, i18n.MsgDistributorTokenModelForbidden, map[string]any{"Model": modelRequest.Model})) + return + } + } - if shouldSelectChannel { - if modelRequest.Model == "" { - abortWithOpenAiMessage(c, http.StatusBadRequest, i18n.T(c, i18n.MsgDistributorModelNameRequired)) + if shouldSelectChannel { + if modelRequest.Model == "" { + abortWithOpenAiMessage(c, http.StatusBadRequest, i18n.T(c, i18n.MsgDistributorModelNameRequired)) + return + } + 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()})) return } - 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()})) + if playgroundRequest.Group != "" { + if !service.GroupInUserUsableGroups(usingGroup, playgroundRequest.Group) && playgroundRequest.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 = playgroundRequest.Group + common.SetContextKey(c, constant.ContextKeyUsingGroup, usingGroup) } + } - // 通道亲和性检查(仅在未选中默认通道时生效) - if channel == nil { - if preferredChannelID, found := service.GetPreferredChannelByAffinity(c, modelRequest.Model, usingGroup); found { - preferred, err := model.CacheGetChannel(preferredChannelID) - if err == nil && preferred != nil && preferred.Status == common.ChannelStatusEnabled { - if usingGroup == "auto" { - userGroup := common.GetContextKeyString(c, constant.ContextKeyUserGroup) - autoGroups := service.GetUserAutoGroup(userGroup) - for _, g := range autoGroups { - if model.IsChannelEnabledForGroupModel(g, modelRequest.Model, preferred.Id) { - selectGroup = g - common.SetContextKey(c, constant.ContextKeyAutoGroup, g) - channel = preferred - service.MarkChannelAffinityUsed(c, g, preferred.Id) - break - } - } - } else if model.IsChannelEnabledForGroupModel(usingGroup, modelRequest.Model, preferred.Id) { + // 通道亲和性检查 + if preferredChannelID, found := service.GetPreferredChannelByAffinity(c, modelRequest.Model, usingGroup); found { + preferred, err := model.CacheGetChannel(preferredChannelID) + if err == nil && preferred != nil && preferred.Status == common.ChannelStatusEnabled { + if usingGroup == "auto" { + userGroup := common.GetContextKeyString(c, constant.ContextKeyUserGroup) + autoGroups := service.GetUserAutoGroup(userGroup) + for _, g := range autoGroups { + if model.IsChannelEnabledForGroupModel(g, modelRequest.Model, preferred.Id) { + selectGroup = g + common.SetContextKey(c, constant.ContextKeyAutoGroup, g) channel = preferred - selectGroup = usingGroup - service.MarkChannelAffinityUsed(c, usingGroup, preferred.Id) + service.MarkChannelAffinityUsed(c, g, preferred.Id) + break } } + } else if model.IsChannelEnabledForGroupModel(usingGroup, modelRequest.Model, preferred.Id) { + channel = preferred + selectGroup = usingGroup + service.MarkChannelAffinityUsed(c, usingGroup, preferred.Id) } } + } - if channel == nil { - channel, selectGroup, err = service.CacheGetRandomSatisfiedChannel(&service.RetryParam{ - Ctx: c, - ModelName: modelRequest.Model, - TokenGroup: usingGroup, - Retry: common.GetPointer(0), - }) - if err != nil { - showGroup := usingGroup - if usingGroup == "auto" { - showGroup = fmt.Sprintf("auto(%s)", selectGroup) - } - message := i18n.T(c, i18n.MsgDistributorGetChannelFailed, map[string]any{"Group": showGroup, "Model": modelRequest.Model, "Error": err.Error()}) - abortWithOpenAiMessage(c, http.StatusServiceUnavailable, message, types.ErrorCodeModelNotFound) - return - } - if channel == nil { - abortWithOpenAiMessage(c, http.StatusServiceUnavailable, i18n.T(c, i18n.MsgDistributorNoAvailableChannel, map[string]any{"Group": usingGroup, "Model": modelRequest.Model}), types.ErrorCodeModelNotFound) - return + if channel == nil { + channel, selectGroup, err = service.CacheGetRandomSatisfiedChannel(&service.RetryParam{ + Ctx: c, + ModelName: modelRequest.Model, + TokenGroup: usingGroup, + Retry: common.GetPointer(0), + }) + if err != nil { + showGroup := usingGroup + if usingGroup == "auto" { + showGroup = fmt.Sprintf("auto(%s)", selectGroup) } + message := i18n.T(c, i18n.MsgDistributorGetChannelFailed, map[string]any{"Group": showGroup, "Model": modelRequest.Model, "Error": err.Error()}) + abortWithOpenAiMessage(c, http.StatusServiceUnavailable, message, types.ErrorCodeModelNotFound) + return + } + if channel == nil { + abortWithOpenAiMessage(c, http.StatusServiceUnavailable, i18n.T(c, i18n.MsgDistributorNoAvailableChannel, map[string]any{"Group": usingGroup, "Model": modelRequest.Model}), types.ErrorCodeModelNotFound) + return } } + } common.SetContextKey(c, constant.ContextKeyRequestStartTime, time.Now()) SetupContextForSelectedChannel(c, channel, modelRequest.Model) c.Next()