diff --git a/controller/relay.go b/controller/relay.go index 9c2d1c4..a9a6179 100644 --- a/controller/relay.go +++ b/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 { diff --git a/relay/helper/matrix_usage_capability.go b/relay/helper/matrix_usage_capability.go new file mode 100644 index 0000000..cc0a56e --- /dev/null +++ b/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 + } +} diff --git a/relay/helper/matrix_usage_capability_test.go b/relay/helper/matrix_usage_capability_test.go new file mode 100644 index 0000000..9524f2a --- /dev/null +++ b/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") +} diff --git a/service/channel_select.go b/service/channel_select.go index a3710ef..8075ef3 100644 --- a/service/channel_select.go +++ b/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 }