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

func ResolveDoubaoAssetChannel(userId int, tokenGroup string) (*model.Channel, error) {
	channel, err := GetBoundDoubaoAssetChannel(userId, tokenGroup)
	if err != nil {
		return nil, err
	}
	if channel != nil {
		return channel, nil
	}

	channel, err = AutoMatchDoubaoAssetChannel(tokenGroup)
	if err != nil {
		return nil, err
	}
	if channel == nil {
		return nil, nil
	}

	if err = model.BindUserAssetChannel(userId, DoubaoAssetChannelType, tokenGroup, channel.Id); err != nil {
		return nil, err
	}
	return channel, nil
}

func GetBoundDoubaoAssetChannel(userId int, tokenGroup string) (*model.Channel, error) {
	binding, err := model.GetUserAssetChannel(userId, DoubaoAssetChannelType, 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 startIdx := 0; ; startIdx += DoubaoAssetChannelPageSize {
		candidates, err := model.GetChannelsByType(startIdx, DoubaoAssetChannelPageSize, true, DoubaoAssetChannelType)
		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 {
			return false
		}
	}
}

func AutoMatchDoubaoAssetChannel(tokenGroup string) (*model.Channel, error) {
	for startIdx := 0; ; startIdx += DoubaoAssetChannelPageSize {
		candidates, err := model.GetChannelsByType(startIdx, DoubaoAssetChannelPageSize, true, DoubaoAssetChannelType)
		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 channel.Type != DoubaoAssetChannelType {
		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 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
}