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},
}

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
	}
}