diff --git a/constant/context_key.go b/constant/context_key.go index 2ba2fe2..2cdce7e 100644 --- a/constant/context_key.go +++ b/constant/context_key.go @@ -19,6 +19,7 @@ const ( ContextKeyTokenModelLimitEnabled ContextKey = "token_model_limit_enabled" ContextKeyTokenModelLimit ContextKey = "token_model_limit" ContextKeyTokenCrossGroupRetry ContextKey = "token_cross_group_retry" + ContextKeyTokenBoundChannelId ContextKey = "token_bound_channel_id" /* channel related keys */ ContextKeyChannelId ContextKey = "channel_id" diff --git a/middleware/auth.go b/middleware/auth.go index 342e7f4..6cb20ab 100644 --- a/middleware/auth.go +++ b/middleware/auth.go @@ -389,6 +389,9 @@ func SetupContextForToken(c *gin.Context, token *model.Token, parts ...string) e } common.SetContextKey(c, constant.ContextKeyTokenGroup, token.Group) common.SetContextKey(c, constant.ContextKeyTokenCrossGroupRetry, token.CrossGroupRetry) + if token.BoundChannelId != nil { + common.SetContextKey(c, constant.ContextKeyTokenBoundChannelId, *token.BoundChannelId) + } if len(parts) > 1 { if model.IsAdmin(token.UserId) { c.Set("specific_channel_id", parts[1]) diff --git a/middleware/distributor.go b/middleware/distributor.go index 9e66cb8..fe08db6 100644 --- a/middleware/distributor.go +++ b/middleware/distributor.go @@ -53,99 +53,127 @@ func Distribute() func(c *gin.Context) { } } else { // Select a channel for the user - // check token model mapping - modelLimitEnable := common.GetContextKeyBool(c, constant.ContextKeyTokenModelLimitEnabled) - if modelLimitEnable { - s, ok := common.GetContextKey(c, constant.ContextKeyTokenModelLimit) + // Check if token has a bound channel + boundChannelId, hasBoundChannel := common.GetContextKey(c, constant.ContextKeyTokenBoundChannelId) + if hasBoundChannel { + id, ok := boundChannelId.(int) if !ok { - // token model limit is empty, all models are not allowed - abortWithOpenAiMessage(c, http.StatusForbidden, i18n.T(c, i18n.MsgDistributorTokenNoModelAccess)) + abortWithOpenAiMessage(c, http.StatusBadRequest, i18n.T(c, i18n.MsgDistributorInvalidChannelId)) 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})) + channel, err = model.GetChannelById(id, true) + if err != nil { + abortWithOpenAiMessage(c, http.StatusBadRequest, i18n.T(c, i18n.MsgDistributorInvalidChannelId)) return } - } - - if shouldSelectChannel { - if modelRequest.Model == "" { - abortWithOpenAiMessage(c, http.StatusBadRequest, i18n.T(c, i18n.MsgDistributorModelNameRequired)) + if channel.Status != common.ChannelStatusEnabled { + abortWithOpenAiMessage(c, http.StatusForbidden, i18n.T(c, i18n.MsgDistributorChannelDisabled)) 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()})) + // Verify the bound channel supports the requested model + if shouldSelectChannel && modelRequest.Model != "" { + usingGroup := common.GetContextKeyString(c, constant.ContextKeyUsingGroup) + if !model.IsChannelEnabledForGroupModel(usingGroup, modelRequest.Model, channel.Id) { + abortWithOpenAiMessage(c, http.StatusServiceUnavailable, i18n.T(c, i18n.MsgDistributorNoAvailableChannel, map[string]any{"Group": usingGroup, "Model": modelRequest.Model}), types.ErrorCodeModelNotFound) 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) + } + } else { + // 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 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 - } + 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 + } + if playgroundRequest.Group != "" { + if !service.GroupInUserUsableGroups(usingGroup, playgroundRequest.Group) && playgroundRequest.Group != usingGroup { + abortWithOpenAiMessage(c, http.StatusForbidden, i18n.T(c, i18n.MsgDistributorGroupAccessDenied)) + return } - } else if model.IsChannelEnabledForGroupModel(usingGroup, modelRequest.Model, preferred.Id) { - channel = preferred - selectGroup = usingGroup - service.MarkChannelAffinityUsed(c, usingGroup, preferred.Id) + usingGroup = playgroundRequest.Group + common.SetContextKey(c, constant.ContextKeyUsingGroup, usingGroup) } } - } - 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) + 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) { + channel = preferred + selectGroup = usingGroup + service.MarkChannelAffinityUsed(c, usingGroup, preferred.Id) + } } - message := i18n.T(c, i18n.MsgDistributorGetChannelFailed, map[string]any{"Group": showGroup, "Model": modelRequest.Model, "Error": err.Error()}) - // 如果错误,但是渠道不为空,说明是数据库一致性问题 - //if channel != nil { - // common.SysError(fmt.Sprintf("渠道不存在:%d", channel.Id)) - // message = "数据库一致性已被破坏,请联系管理员" - //} - 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 + 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()}) + // 如果错误,但是渠道不为空,说明是数据库一致性问题 + //if channel != nil { + // common.SysError(fmt.Sprintf("渠道不存在:%d", channel.Id)) + // message = "数据库一致性已被破坏,请联系管理员" + //} + 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 + } } } }