feat(task-pricing): gate channels for matrix usage billing

controller/relay.go

@@ -191,10 +191,11 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) {
 	}()
 
 	retryParam := &service.RetryParam{
-		Ctx:        c,
-		TokenGroup: relayInfo.TokenGroup,
-		ModelName:  relayInfo.OriginModelName,
-		Retry:      common.GetPointer(0),
+		Ctx:                       c,
+		TokenGroup:                relayInfo.TokenGroup,
+		ModelName:                 relayInfo.OriginModelName,
+		Retry:                     common.GetPointer(0),
+		RequireMatrixUsageBilling: relayInfo.RequireMatrixUsageBilling,
 	}
 
 	for ; retryParam.GetRetry() <= common.RetryTimes; retryParam.IncreaseRetry() {
@@ -314,6 +315,9 @@ func getChannel(c *gin.Context, info *relaycommon.RelayInfo, retryParam *service
 	info.PriceData.GroupRatioInfo = helper.HandleGroupRatio(c, info)
 
 	if err != nil {
+		if apiErr, ok := err.(*types.NewAPIError); ok {
+			return nil, apiErr
+		}
 		return nil, types.NewError(fmt.Errorf("获取分组 %s 下模型 %s 的可用渠道失败（retry）: %s", selectGroup, info.OriginModelName, err.Error()), types.ErrorCodeGetChannelFailed, types.ErrOptionWithSkipRetry())
 	}
 	if channel == nil {

relay/helper/matrix_usage_capability.go

@@ -0,0 +1,60 @@
+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
+	}
+}

relay/helper/matrix_usage_capability_test.go

@@ -0,0 +1,27 @@
+package helper
+
+import (
+	"testing"
+
+	"github.com/QuantumNous/new-api/constant"
+	"github.com/stretchr/testify/require"
+)
+
+func TestSupportsMatrixUsageBilling(t *testing.T) {
+	require.True(t, SupportsMatrixUsageBilling("seedance-2", constant.ChannelTypeDoubaoVideo, "seedance-2"))
+	require.False(t, SupportsMatrixUsageBilling("seedance-2", constant.ChannelTypeOpenAI, "seedance-2"))
+	require.False(t, SupportsMatrixUsageBilling("unknown", constant.ChannelTypeDoubaoVideo, "unknown"))
+}
+
+func TestResolvePricingModelMapping(t *testing.T) {
+	mapped, err := ResolvePricingModelMapping("seedance-2", `{"seedance-2":"doubao-video-seedance"}`)
+	require.NoError(t, err)
+	require.Equal(t, "doubao-video-seedance", mapped)
+
+	mapped, err = ResolvePricingModelMapping("seedance-2", `{"seedance-2":"a","a":"b"}`)
+	require.NoError(t, err)
+	require.Equal(t, "b", mapped)
+
+	_, err = ResolvePricingModelMapping("seedance-2", `{"seedance-2":"a","a":"seedance-2"}`)
+	require.ErrorContains(t, err, "model_mapping_contains_cycle")
+}

service/channel_select.go

@@ -2,21 +2,26 @@ package service
 
 import (
 	"errors"
+	"fmt"
+	"net/http"
 
 	"github.com/QuantumNous/new-api/common"
 	"github.com/QuantumNous/new-api/constant"
 	"github.com/QuantumNous/new-api/logger"
 	"github.com/QuantumNous/new-api/model"
+	"github.com/QuantumNous/new-api/relay/helper"
 	"github.com/QuantumNous/new-api/setting"
+	"github.com/QuantumNous/new-api/types"
 	"github.com/gin-gonic/gin"
 )
 
 type RetryParam struct {
-	Ctx          *gin.Context
-	TokenGroup   string
-	ModelName    string
-	Retry        *int
-	resetNextTry bool
+	Ctx                       *gin.Context
+	TokenGroup                string
+	ModelName                 string
+	Retry                     *int
+	RequireMatrixUsageBilling bool
+	resetNextTry              bool
 }
 
 func (p *RetryParam) GetRetry() int {
@@ -158,5 +163,19 @@ func CacheGetRandomSatisfiedChannel(param *RetryParam) (*model.Channel, string,
 			return nil, param.TokenGroup, err
 		}
 	}
+	if channel != nil && param.RequireMatrixUsageBilling {
+		mappedModel, mapErr := helper.ResolvePricingModelMapping(param.ModelName, channel.GetModelMapping())
+		if mapErr != nil {
+			return nil, selectGroup, mapErr
+		}
+		if !helper.SupportsMatrixUsageBilling(param.ModelName, channel.Type, mappedModel) {
+			return nil, selectGroup, types.NewErrorWithStatusCode(
+				fmt.Errorf("matrix_usage_billing_channel_unavailable"),
+				types.ErrorCode("matrix_usage_billing_channel_unavailable"),
+				http.StatusBadRequest,
+				types.ErrOptionWithSkipRetry(),
+			)
+		}
+	}
 	return channel, selectGroup, nil
 }