new-api/service/doubao_asset_channel.go

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
}