new-api/service/asset_adapter.go

package service

import (
	"context"
	"net/http"
	"sort"
	"strings"

	"github.com/QuantumNous/new-api/constant"
	"github.com/QuantumNous/new-api/model"
)

const (
	AssetErrorInvalidRequest        = "invalid_request_error"
	AssetErrorChannelNotFound       = "asset_channel_not_found"
	AssetErrorBindingInvalid        = "asset_channel_binding_invalid"
	AssetErrorOperationNotSupported = "asset_operation_not_supported"
	AssetErrorNotFound              = "asset_not_found"
	AssetErrorUpstream              = "upstream_error"
	AssetErrorServer                = "server_error"
)

type AssetError struct {
	Type       string
	Message    string
	HTTPStatus int
}

func newAssetError(errType string, message string, status int) *AssetError {
	if status == 0 {
		status = http.StatusBadGateway
	}
	return &AssetError{Type: errType, Message: message, HTTPStatus: status}
}

type AssetRequest struct {
	Action  AssetActionSpec
	Version string
	Body    map[string]any
	RawBody []byte
}

type AssetUpstreamResponse struct {
	StatusCode int
	Header     http.Header
	Body       []byte
}

type AssetAdapter interface {
	Name() string
	Supports(operation AssetOperation) bool
	DoAssetRequest(ctx context.Context, channel *model.Channel, req AssetRequest) (*AssetUpstreamResponse, *AssetError)
}

var assetAdapters = map[int]AssetAdapter{}

func init() {
	registerDefaultAssetAdapters()
}

func registerDefaultAssetAdapters() {
	assetAdapters[constant.ChannelTypeChinaMobileSeedance] = NewChinaMobileAssetAdapter()
	assetAdapters[constant.ChannelTypeDoubaoVideoCompatibleAiping] = NewCompatibleAssetAdapter("aiping_asset", []AssetOperation{
		AssetOperationAssetCreate,
		AssetOperationAssetList,
		AssetOperationAssetGet,
		AssetOperationAssetUpdate,
		AssetOperationAssetDelete,
	})
	assetAdapters[constant.ChannelTypeDoubaoVideoCompatibleTianyiYun] = NewCompatibleAssetAdapter("tianyiyun_asset", []AssetOperation{
		AssetOperationAssetCreate,
		AssetOperationAssetList,
		AssetOperationAssetGet,
		AssetOperationAssetUpdate,
		AssetOperationAssetDelete,
	})
}

func GetAssetAdapter(channelType int) (AssetAdapter, bool) {
	adapter, ok := assetAdapters[channelType]
	return adapter, ok
}

func OverrideAssetAdapterForTest(channelType int, adapter AssetAdapter) func() {
	oldAdapter, hadOldAdapter := assetAdapters[channelType]
	assetAdapters[channelType] = adapter
	return func() {
		if hadOldAdapter {
			assetAdapters[channelType] = oldAdapter
		} else {
			delete(assetAdapters, channelType)
		}
	}
}

func RegisteredAssetChannelTypes() []int {
	types := make([]int, 0, len(assetAdapters))
	for channelType := range assetAdapters {
		types = append(types, channelType)
	}
	sort.Ints(types)
	return types
}

func assetChannelHasKey(channel *model.Channel) bool {
	return channel != nil && strings.TrimSpace(channel.Key) != "" && len(channel.GetKeys()) > 0
}