new-api/relay/helper/matrix_usage_capability.go

package helper

import (
	"errors"

	commonpkg "github.com/QuantumNous/new-api/common"
	"github.com/QuantumNous/new-api/constant"
)

type MatrixUsageCapability struct {
	BillingModelName  string
	UpstreamModelName string
	ChannelType       int
}

var matrixUsageCapabilities = []MatrixUsageCapability{
	{BillingModelName: "seedance-2", UpstreamModelName: "seedance-2", ChannelType: constant.ChannelTypeDoubaoVideo},
	{BillingModelName: "doubao-video", UpstreamModelName: "doubao-video", ChannelType: constant.ChannelTypeDoubaoVideo},
	{BillingModelName: "doubao-video-seedance", UpstreamModelName: "doubao-video-seedance", ChannelType: constant.ChannelTypeDoubaoVideo},
	{BillingModelName: "doubao-seedance-2-0-260128", UpstreamModelName: "doubao-seedance-2-0-260128", ChannelType: constant.ChannelTypeDoubaoVideo},
	{BillingModelName: "doubao-seedance-2-0-fast-260128", UpstreamModelName: "doubao-seedance-2-0-fast-260128", ChannelType: constant.ChannelTypeDoubaoVideo},
	{BillingModelName: "doubao-seedance-2-0-260128", UpstreamModelName: "doubao-seedance-2-0-260128", ChannelType: constant.ChannelTypeDoubaoVideoCompatibleAiping},
	{BillingModelName: "doubao-seedance-2-0-fast-260128", UpstreamModelName: "doubao-seedance-2-0-fast-260128", ChannelType: constant.ChannelTypeDoubaoVideoCompatibleAiping},
	{BillingModelName: "cdance2.0-0611", UpstreamModelName: "cdance2.0-0611", ChannelType: constant.ChannelTypeDoubaoVideoCompatibleTianyiYun},
	{BillingModelName: "cdance2.0-fast-0611", UpstreamModelName: "cdance2.0-fast-0611", ChannelType: constant.ChannelTypeDoubaoVideoCompatibleTianyiYun},
	// Keep in sync with SupportsAnyMatrixUsageBillingModel in setting/ratio_setting/model_pricing.go
	{BillingModelName: "kling-v1", UpstreamModelName: "kling-v1", ChannelType: constant.ChannelTypeKlingAiping},
	{BillingModelName: "kling-v1", UpstreamModelName: "Kling-V1", ChannelType: constant.ChannelTypeKlingAiping},
	{BillingModelName: "kling-v1-6", UpstreamModelName: "kling-v1-6", ChannelType: constant.ChannelTypeKlingAiping},
	{BillingModelName: "kling-v1-6", UpstreamModelName: "Kling-V1.6", ChannelType: constant.ChannelTypeKlingAiping},
	{BillingModelName: "kling-v2-6", UpstreamModelName: "kling-v2-6", ChannelType: constant.ChannelTypeKlingAiping},
	{BillingModelName: "kling-v2-6", UpstreamModelName: "Kling-V2.6", ChannelType: constant.ChannelTypeKlingAiping},
	{BillingModelName: "kling-v3", UpstreamModelName: "kling-v3", ChannelType: constant.ChannelTypeKlingAiping},
	{BillingModelName: "kling-v3", UpstreamModelName: "Kling-V3", ChannelType: constant.ChannelTypeKlingAiping},
	{BillingModelName: "kling-video-o1", UpstreamModelName: "kling-video-o1", ChannelType: constant.ChannelTypeKlingAiping},
	{BillingModelName: "kling-video-o1", UpstreamModelName: "Kling-Video-O1", ChannelType: constant.ChannelTypeKlingAiping},
	{BillingModelName: "kling-v3-omni", UpstreamModelName: "kling-v3-omni", ChannelType: constant.ChannelTypeKlingAiping},
	{BillingModelName: "kling-v3-omni", UpstreamModelName: "Kling-V3-Omni", ChannelType: constant.ChannelTypeKlingAiping},
}

func SupportsMatrixUsageBilling(modelName string, channelType int, mappedModel string) bool {
	if mappedModel == "" {
		mappedModel = modelName
	}
	for _, capability := range matrixUsageCapabilities {
		if capability.BillingModelName == modelName &&
			capability.UpstreamModelName == mappedModel &&
			capability.ChannelType == channelType {
			return true
		}
	}
	return false
}

func ResolvePricingModelMapping(modelName string, modelMapping string) (string, error) {
	if modelMapping == "" || modelMapping == "{}" {
		return modelName, nil
	}
	modelMap := map[string]string{}
	if err := commonpkg.Unmarshal([]byte(modelMapping), &modelMap); err != nil {
		return "", err
	}
	current := modelName
	visited := map[string]bool{current: true}
	for {
		next := modelMap[current]
		if next == "" {
			return current, nil
		}
		if visited[next] {
			if next == current {
				return current, nil
			}
			return "", errors.New("model_mapping_contains_cycle")
		}
		visited[next] = true
		current = next
	}
}