package service import ( "errors" "strings" "github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/constant" "github.com/QuantumNous/new-api/model" "gorm.io/gorm" ) const DoubaoAssetChannelPageSize = 50 const DoubaoAssetChannelType = constant.ChannelTypeDoubaoVideoCompatibleAiping var DoubaoAssetChannelTypes = []int{ constant.ChannelTypeDoubaoVideoCompatibleAiping, constant.ChannelTypeDoubaoVideoCompatibleTianyiYun, } func ResolveDoubaoAssetChannel(userId int, tokenGroup string) (*model.Channel, error) { return ResolveDoubaoAssetChannelForModel(userId, tokenGroup, "") } func ResolveDoubaoAssetChannelForModel(userId int, tokenGroup string, modelName string) (*model.Channel, error) { channelType := DoubaoAssetChannelTypeForModel(modelName) channel, err := getBoundDoubaoAssetChannelByType(userId, tokenGroup, channelType) if err != nil { return nil, err } if channel != nil { return channel, nil } channel, err = autoMatchDoubaoAssetChannelByType(tokenGroup, channelType) if err != nil { return nil, err } if channel == nil { return nil, nil } if err = model.BindUserAssetChannel(userId, channel.Type, tokenGroup, channel.Id); err != nil { return nil, err } return channel, nil } func GetBoundDoubaoAssetChannel(userId int, tokenGroup string) (*model.Channel, error) { return getBoundDoubaoAssetChannelByType(userId, tokenGroup, DoubaoAssetChannelType) } func GetBoundDoubaoAssetChannelForModel(userId int, tokenGroup string, modelName string) (*model.Channel, error) { return getBoundDoubaoAssetChannelByType(userId, tokenGroup, DoubaoAssetChannelTypeForModel(modelName)) } func getBoundDoubaoAssetChannelByType(userId int, tokenGroup string, channelType int) (*model.Channel, error) { binding, err := model.GetUserAssetChannel(userId, channelType, tokenGroup) if err != nil { return nil, err } if binding == nil { return nil, nil } channel, err := model.CacheGetChannel(binding.ChannelId) if err != nil { if errors.Is(err, gorm.ErrRecordNotFound) || common.MemoryCacheEnabled { return nil, nil } return nil, err } if !IsUsableDoubaoAssetChannel(channel, tokenGroup) { return nil, nil } return channel, nil } func HasDoubaoVideoChannelForGroupModel(tokenGroup string, modelName string) bool { if tokenGroup == "" || modelName == "" { return false } for _, channelType := range DoubaoAssetChannelTypes { for startIdx := 0; ; startIdx += DoubaoAssetChannelPageSize { candidates, err := model.GetChannelsByType(startIdx, DoubaoAssetChannelPageSize, true, channelType) if err != nil { return false } for _, candidate := range candidates { if candidate == nil || !MatchDoubaoAssetGroup(candidate.GetGroups(), tokenGroup) { continue } if model.IsChannelEnabledForGroupModel(tokenGroup, modelName, candidate.Id) { return true } } if len(candidates) < DoubaoAssetChannelPageSize { break } } } return false } func AutoMatchDoubaoAssetChannel(tokenGroup string) (*model.Channel, error) { return autoMatchDoubaoAssetChannelByType(tokenGroup, DoubaoAssetChannelType) } func autoMatchDoubaoAssetChannelByType(tokenGroup string, channelType int) (*model.Channel, error) { for startIdx := 0; ; startIdx += DoubaoAssetChannelPageSize { candidates, err := model.GetChannelsByType(startIdx, DoubaoAssetChannelPageSize, true, channelType) if err != nil { return nil, err } for _, candidate := range candidates { if candidate == nil || !MatchDoubaoAssetGroup(candidate.GetGroups(), tokenGroup) { continue } channel, err := model.CacheGetChannel(candidate.Id) if err != nil { if errors.Is(err, gorm.ErrRecordNotFound) || common.MemoryCacheEnabled { continue } return nil, err } if IsUsableDoubaoAssetChannel(channel, tokenGroup) { return channel, nil } } if len(candidates) < DoubaoAssetChannelPageSize { return nil, nil } } } func IsUsableDoubaoAssetChannel(channel *model.Channel, tokenGroup string) bool { if channel == nil { return false } if !IsDoubaoAssetChannelType(channel.Type) { return false } if channel.Status != common.ChannelStatusEnabled { return false } if strings.TrimSpace(channel.Key) == "" || len(channel.GetKeys()) == 0 { return false } return MatchDoubaoAssetGroup(channel.GetGroups(), tokenGroup) } func IsDoubaoAssetChannelType(channelType int) bool { for _, t := range DoubaoAssetChannelTypes { if channelType == t { return true } } return false } func DoubaoAssetChannelTypeForModel(modelName string) int { if strings.HasPrefix(strings.ToLower(strings.TrimSpace(modelName)), "cdance") { return constant.ChannelTypeDoubaoVideoCompatibleTianyiYun } return DoubaoAssetChannelType } func MatchDoubaoAssetGroup(channelGroups []string, tokenGroup string) bool { tokenGroup = strings.TrimSpace(tokenGroup) if tokenGroup == "" { return false } for _, group := range channelGroups { if strings.TrimSpace(group) == tokenGroup { return true } } return false }