new-api/service/video_asset_channel.go

package service

import (
	"errors"
	"sort"
	"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 VideoAssetFamilies() []VideoAssetFamily {
	return []VideoAssetFamily{VideoAssetFamilySeedance, VideoAssetFamilyKling}
}

func VideoAssetChannelTypesForFamily(family VideoAssetFamily) []int {
	switch family {
	case VideoAssetFamilySeedance:
		// DoubaoVideo (official Volcengine Ark) joins the family so a user's
		// asset upload channel and video task channel stay the same; assets
		// are referenced by asset:// ids passed through to the upstream.
		return []int{
			constant.ChannelTypeDoubaoVideo,
			constant.ChannelTypeDoubaoVideoCompatibleAiping,
			constant.ChannelTypeDoubaoVideoCompatibleTianyiYun,
			constant.ChannelTypeChinaMobileSeedance,
		}
	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 ClearVideoAssetChannelBinding(userId int, tokenGroup string, family VideoAssetFamily) error {
	channelTypes := VideoAssetChannelTypesForFamily(family)
	if len(channelTypes) == 0 {
		return nil
	}
	return model.DB.Transaction(func(tx *gorm.DB) error {
		return model.DeleteUserAssetChannelsByTypesWithTx(tx, userId, channelTypes, tokenGroup)
	})
}

func GetVideoAssetChannelCandidates(tokenGroup string, family VideoAssetFamily) ([]*model.Channel, error) {
	channelTypes := VideoAssetChannelTypesForFamily(family)
	if len(channelTypes) == 0 {
		return nil, nil
	}
	candidates := make(map[int]*model.Channel)
	for _, channelType := range channelTypes {
		for startIdx := 0; ; startIdx += DoubaoAssetChannelPageSize {
			channels, err := model.GetChannelsByType(startIdx, DoubaoAssetChannelPageSize, true, channelType)
			if err != nil {
				return nil, err
			}
			for _, candidate := range channels {
				if candidate == nil || candidate.Status != common.ChannelStatusEnabled || !MatchDoubaoAssetGroup(candidate.GetGroups(), tokenGroup) {
					continue
				}
				channel, err := model.GetChannelById(candidate.Id, true)
				if err != nil {
					return nil, err
				}
				if channel.Status == common.ChannelStatusEnabled && strings.TrimSpace(channel.Key) != "" && MatchDoubaoAssetGroup(channel.GetGroups(), tokenGroup) {
					channel.Key = ""
					channel.Keys = nil
					candidates[channel.Id] = channel
				}
			}
			if len(channels) < DoubaoAssetChannelPageSize {
				break
			}
		}
	}
	result := make([]*model.Channel, 0, len(candidates))
	for _, channel := range candidates {
		result = append(result, channel)
	}
	sort.Slice(result, func(i, j int) bool { return result[i].Id < result[j].Id })
	return result, nil
}

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
	}
	// Official Volcengine channels additionally need the asset AK/SK
	// credential before they can serve the asset family; without it they are
	// skipped so binding/auto-match falls through to compatible channels.
	if !assetCredentialConfigured(channel) {
		return false
	}
	modelName = strings.TrimSpace(modelName)
	if modelName == "" {
		return true
	}
	return model.IsChannelEnabledForGroupModel(tokenGroup, modelName, channel.Id)
}