new-api/service/kling_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 KlingAssetChannelType = constant.ChannelTypeKlingAiping

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

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

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

func GetBoundKlingAssetChannel(userId int, tokenGroup string) (*model.Channel, error) {
	binding, err := model.GetUserAssetChannel(userId, KlingAssetChannelType, 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 !IsUsableKlingAssetChannel(channel, tokenGroup) {
		return nil, nil
	}
	return channel, nil
}

func AutoMatchKlingAssetChannel(tokenGroup string) (*model.Channel, error) {
	for startIdx := 0; ; startIdx += DoubaoAssetChannelPageSize {
		candidates, err := model.GetChannelsByType(startIdx, DoubaoAssetChannelPageSize, true, KlingAssetChannelType)
		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 IsUsableKlingAssetChannel(channel, tokenGroup) {
				return channel, nil
			}
		}
		if len(candidates) < DoubaoAssetChannelPageSize {
			return nil, nil
		}
	}
}

func IsUsableKlingAssetChannel(channel *model.Channel, tokenGroup string) bool {
	if channel == nil {
		return false
	}
	if channel.Type != KlingAssetChannelType {
		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)
}