new-api/service/asset_resolver.go

package service

import (
	"errors"
	"fmt"
	"net/http"

	"github.com/QuantumNous/new-api/common"
	"github.com/QuantumNous/new-api/model"
	"gorm.io/gorm"
)

func ResolveAssetChannelForOperation(userID int, tokenGroup string, operation AssetOperation) (*model.Channel, AssetAdapter, *AssetError) {
	bindings, err := model.GetUserAssetChannelsByTypes(userID, RegisteredAssetChannelTypes(), tokenGroup)
	if err != nil {
		return nil, nil, newAssetError(AssetErrorServer, err.Error(), http.StatusInternalServerError)
	}
	if len(bindings) > 0 {
		binding := bindings[0]
		channel, err := model.CacheGetChannel(binding.ChannelId)
		if err != nil {
			if errors.Is(err, gorm.ErrRecordNotFound) || common.MemoryCacheEnabled {
				return nil, nil, newAssetError(AssetErrorBindingInvalid, "bound asset channel not found", http.StatusBadGateway)
			}
			return nil, nil, newAssetError(AssetErrorServer, err.Error(), http.StatusInternalServerError)
		}
		adapter, ok := GetAssetAdapter(channel.Type)
		if channel.Status != common.ChannelStatusEnabled || !assetChannelHasKey(channel) || !MatchDoubaoAssetGroup(channel.GetGroups(), tokenGroup) || !ok {
			return nil, nil, newAssetError(AssetErrorBindingInvalid, "bound asset channel is not available for asset library", http.StatusBadGateway)
		}
		if !adapter.Supports(operation) {
			return nil, nil, newAssetError(AssetErrorOperationNotSupported, fmt.Sprintf("asset operation %s is not supported by bound channel", operation), http.StatusBadRequest)
		}
		return channel, adapter, nil
	}

	channel, adapter, err := autoMatchAssetChannelForOperation(tokenGroup, operation)
	if err != nil {
		return nil, nil, newAssetError(AssetErrorServer, err.Error(), http.StatusInternalServerError)
	}
	if channel == nil {
		return nil, nil, newAssetError(AssetErrorChannelNotFound, "no available asset channel supports requested operation", http.StatusBadGateway)
	}
	return channel, adapter, nil
}

func autoMatchAssetChannelForOperation(tokenGroup string, operation AssetOperation) (*model.Channel, AssetAdapter, error) {
	for _, channelType := range RegisteredAssetChannelTypes() {
		adapter, ok := GetAssetAdapter(channelType)
		if !ok || !adapter.Supports(operation) {
			continue
		}
		for startIdx := 0; ; startIdx += DoubaoAssetChannelPageSize {
			candidates, err := model.GetChannelsByType(startIdx, DoubaoAssetChannelPageSize, true, channelType)
			if err != nil {
				return nil, 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, nil, err
				}
				if channel.Status == common.ChannelStatusEnabled && assetChannelHasKey(channel) && MatchDoubaoAssetGroup(channel.GetGroups(), tokenGroup) {
					return channel, adapter, nil
				}
			}
			if len(candidates) < DoubaoAssetChannelPageSize {
				break
			}
		}
	}
	return nil, nil, nil
}