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" ) type VideoAssetFamily string const ( VideoAssetFamilySeedance VideoAssetFamily = "seedance" VideoAssetFamilyKling VideoAssetFamily = "kling" ) func VideoAssetChannelTypesForFamily(family VideoAssetFamily) []int { switch family { case VideoAssetFamilySeedance: return []int{ constant.ChannelTypeDoubaoVideoCompatibleAiping, constant.ChannelTypeDoubaoVideoCompatibleTianyiYun, } case VideoAssetFamilyKling: return []int{ constant.ChannelTypeKlingAiping, } default: return nil } } func VideoAssetFamilyForChannelType(channelType int) (VideoAssetFamily, bool) { for _, family := range []VideoAssetFamily{VideoAssetFamilySeedance, VideoAssetFamilyKling} { if channelTypeAllowed(channelType, VideoAssetChannelTypesForFamily(family)) { return family, true } } return "", false } func GetBoundVideoAssetChannelForModel(userId int, tokenGroup string, modelName string, family VideoAssetFamily) (*model.Channel, error) { bindings, err := model.GetUserAssetChannelsByTypes(userId, VideoAssetChannelTypesForFamily(family), tokenGroup) if err != nil { return nil, err } for _, binding := range bindings { channel, err := model.CacheGetChannel(binding.ChannelId) if err != nil { if errors.Is(err, gorm.ErrRecordNotFound) || common.MemoryCacheEnabled { continue } return nil, err } if IsUsableVideoAssetChannelForFamily(channel, tokenGroup, modelName, family) { return channel, nil } } return nil, nil } func ResolveVideoAssetChannelForModel(userId int, tokenGroup string, modelName string, family VideoAssetFamily) (*model.Channel, error) { channel, err := GetBoundVideoAssetChannelForModel(userId, tokenGroup, modelName, family) if err != nil { return nil, err } if channel != nil { return channel, nil } channel, err = autoMatchVideoAssetChannelForModel(tokenGroup, modelName, family) if err != nil { return nil, err } if channel == nil { return nil, nil } if err = BindVideoAssetChannel(userId, tokenGroup, channel, family); err != nil { return nil, err } return channel, nil } func BindVideoAssetChannel(userId int, tokenGroup string, channel *model.Channel, family VideoAssetFamily) error { if channel == nil { return nil } channelTypes := VideoAssetChannelTypesForFamily(family) if len(channelTypes) == 0 || !channelTypeAllowed(channel.Type, channelTypes) { return nil } otherTypes := make([]int, 0, len(channelTypes)) for _, channelType := range channelTypes { if channelType != channel.Type { otherTypes = append(otherTypes, channelType) } } return model.DB.Transaction(func(tx *gorm.DB) error { if err := model.DeleteUserAssetChannelsByTypesWithTx(tx, userId, otherTypes, tokenGroup); err != nil { return err } return model.BindUserAssetChannelWithTx(tx, userId, channel.Type, tokenGroup, channel.Id) }) } func HasVideoAssetChannelForGroupModel(tokenGroup string, modelName string, family VideoAssetFamily) bool { tokenGroup = strings.TrimSpace(tokenGroup) modelName = strings.TrimSpace(modelName) if tokenGroup == "" || modelName == "" { return false } channel, err := autoMatchVideoAssetChannelForModel(tokenGroup, modelName, family) return err == nil && channel != nil } func autoMatchVideoAssetChannelForModel(tokenGroup string, modelName string, family VideoAssetFamily) (*model.Channel, error) { for _, channelType := range VideoAssetChannelTypesForFamily(family) { 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 IsUsableVideoAssetChannelForFamily(channel, tokenGroup, modelName, family) { return channel, nil } } if len(candidates) < DoubaoAssetChannelPageSize { break } } } return nil, nil } func IsUsableVideoAssetChannelForFamily(channel *model.Channel, tokenGroup string, modelName string, family VideoAssetFamily) bool { if channel == nil { return false } channelTypes := VideoAssetChannelTypesForFamily(family) if len(channelTypes) == 0 || !channelTypeAllowed(channel.Type, channelTypes) { return false } if channel.Status != common.ChannelStatusEnabled { return false } if strings.TrimSpace(channel.Key) == "" || len(channel.GetKeys()) == 0 { return false } if !MatchDoubaoAssetGroup(channel.GetGroups(), tokenGroup) { return false } modelName = strings.TrimSpace(modelName) if modelName == "" { return true } return model.IsChannelEnabledForGroupModel(tokenGroup, modelName, channel.Id) }