From fc1259f583664668ff27809a6acfe1129e361272 Mon Sep 17 00:00:00 2001 From: CaIon Date: Tue, 7 Jul 2026 21:12:13 +0800 Subject: [PATCH] refactor(price): improve handling of other ratios in PriceData --- controller/relay.go | 2 +- relay/chat_completions_via_responses_test.go | 45 +++++++++ relay/image_handler.go | 2 +- relay/relay_task.go | 40 +++----- service/task_billing.go | 28 ++++-- service/task_billing_test.go | 99 ++++++++++++++++++++ service/text_quota.go | 13 +-- types/price_data.go | 83 ++++++++++++++-- 8 files changed, 256 insertions(+), 56 deletions(-) diff --git a/controller/relay.go b/controller/relay.go index ee24100d..6e91ccb6 100644 --- a/controller/relay.go +++ b/controller/relay.go @@ -588,7 +588,7 @@ func RelayTask(c *gin.Context) { ModelPrice: relayInfo.PriceData.ModelPrice, GroupRatio: relayInfo.PriceData.GroupRatioInfo.GroupRatio, ModelRatio: relayInfo.PriceData.ModelRatio, - OtherRatios: relayInfo.PriceData.OtherRatios, + OtherRatios: relayInfo.PriceData.OtherRatios(), OriginModelName: relayInfo.OriginModelName, PerCallBilling: common.StringsContains(constant.TaskPricePatches, relayInfo.OriginModelName) || relayInfo.PriceData.UsePrice, } diff --git a/relay/chat_completions_via_responses_test.go b/relay/chat_completions_via_responses_test.go index ae96d0fd..18587874 100644 --- a/relay/chat_completions_via_responses_test.go +++ b/relay/chat_completions_via_responses_test.go @@ -1,9 +1,13 @@ package relay import ( + "math" "testing" + relaycommon "github.com/QuantumNous/new-api/relay/common" + "github.com/QuantumNous/new-api/types" "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" ) func TestIsResponsesEventStreamContentType(t *testing.T) { @@ -24,3 +28,44 @@ func TestIsResponsesEventStreamContentType(t *testing.T) { }) } } + +func TestRecalcQuotaFromRatiosIgnoresInvalidMultipliers(t *testing.T) { + info := &relaycommon.RelayInfo{ + PriceData: types.PriceData{ + Quota: 100, + }, + } + info.PriceData.AddOtherRatio("duration", 2) + + quota, ok := recalcQuotaFromRatios(info, map[string]float64{ + "duration": 3, + "zero": 0, + "negative": -1, + "nan": math.NaN(), + "inf": math.Inf(1), + }) + + require.True(t, ok) + assert.Equal(t, 150, quota) + assert.True(t, info.PriceData.HasOtherRatio("duration")) +} + +func TestRecalcQuotaFromRatiosRejectsAllInvalidAdjustedRatios(t *testing.T) { + info := &relaycommon.RelayInfo{ + PriceData: types.PriceData{ + Quota: 100, + }, + } + info.PriceData.AddOtherRatio("duration", 2) + + quota, ok := recalcQuotaFromRatios(info, map[string]float64{ + "zero": 0, + "negative": -1, + "nan": math.NaN(), + "inf": math.Inf(1), + }) + + require.False(t, ok) + assert.Equal(t, 0, quota) + assert.True(t, info.PriceData.HasOtherRatio("duration")) +} diff --git a/relay/image_handler.go b/relay/image_handler.go index aa76f197..52aa3e78 100644 --- a/relay/image_handler.go +++ b/relay/image_handler.go @@ -128,7 +128,7 @@ func ImageHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *type // Adaptors may have already set a more accurate count from the // upstream response; only set the default when they haven't. if info.PriceData.UsePrice { // only price model use N ratio - if _, hasN := info.PriceData.OtherRatios["n"]; !hasN { + if !info.PriceData.HasOtherRatio("n") { info.PriceData.AddOtherRatio("n", float64(imageN)) } } diff --git a/relay/relay_task.go b/relay/relay_task.go index c103f7d5..fb384d18 100644 --- a/relay/relay_task.go +++ b/relay/relay_task.go @@ -196,12 +196,7 @@ func RelayTaskSubmit(c *gin.Context, info *relaycommon.RelayInfo) (*TaskSubmitRe // 6. 将 OtherRatios 应用到基础额度（饱和转换，防止溢出成负数） if !common.StringsContains(constant.TaskPricePatches, modelName) { - quotaWithRatios := float64(info.PriceData.Quota) - for _, ra := range info.PriceData.OtherRatios { - if ra != 1.0 { - quotaWithRatios *= ra - } - } + quotaWithRatios := info.PriceData.ApplyOtherRatiosToFloat(float64(info.PriceData.Quota)) quota, clamp := common.QuotaFromFloatChecked(quotaWithRatios) info.PriceData.Quota = quota noteTaskQuotaClamp(info, clamp) @@ -232,7 +227,7 @@ func RelayTaskSubmit(c *gin.Context, info *relaycommon.RelayInfo) (*TaskSubmitRe } // 10. 返回 OtherRatios 给下游（header 必须在 DoResponse 写 body 之前设置） - otherRatios := info.PriceData.OtherRatios + otherRatios := info.PriceData.OtherRatios() if otherRatios == nil { otherRatios = map[string]float64{} } @@ -248,10 +243,12 @@ func RelayTaskSubmit(c *gin.Context, info *relaycommon.RelayInfo) (*TaskSubmitRe // 11. 提交后计费调整：让适配器根据上游实际返回调整 OtherRatios finalQuota := info.PriceData.Quota if adjustedRatios := adaptor.AdjustBillingOnSubmit(info, taskData); len(adjustedRatios) > 0 { - // 基于调整后的 ratios 重新计算 quota - finalQuota = recalcQuotaFromRatios(info, adjustedRatios) - info.PriceData.OtherRatios = adjustedRatios - info.PriceData.Quota = finalQuota + if adjustedQuota, ok := recalcQuotaFromRatios(info, adjustedRatios); ok { + // 基于调整后的 ratios 重新计算 quota + finalQuota = adjustedQuota + info.PriceData.ReplaceOtherRatios(adjustedRatios) + info.PriceData.Quota = finalQuota + } } return &TaskSubmitResult{ @@ -264,25 +261,18 @@ func RelayTaskSubmit(c *gin.Context, info *relaycommon.RelayInfo) (*TaskSubmitRe // recalcQuotaFromRatios 根据 adjustedRatios 重新计算 quota。 // 公式: baseQuota × ∏(ratio) — 其中 baseQuota 是不含 OtherRatios 的基础额度。 -func recalcQuotaFromRatios(info *relaycommon.RelayInfo, ratios map[string]float64) int { +func recalcQuotaFromRatios(info *relaycommon.RelayInfo, ratios map[string]float64) (int, bool) { // 从 PriceData 获取不含 OtherRatios 的基础价格 - baseQuota := float64(info.PriceData.Quota) - // 先除掉原有的 OtherRatios 恢复基础额度 - for _, ra := range info.PriceData.OtherRatios { - if ra != 1.0 && ra > 0 { - baseQuota /= ra - } + baseQuota := info.PriceData.RemoveOtherRatiosFromFloat(float64(info.PriceData.Quota)) + priceData := info.PriceData + if !priceData.ReplaceOtherRatios(ratios) { + return 0, false } // 应用新的 ratios - result := baseQuota - for _, ra := range ratios { - if ra != 1.0 { - result *= ra - } - } + result := priceData.ApplyOtherRatiosToFloat(baseQuota) quota, clamp := common.QuotaFromFloatChecked(result) noteTaskQuotaClamp(info, clamp) - return quota + return quota, true } // noteTaskQuotaClamp records the first quota saturation event onto the task's diff --git a/service/task_billing.go b/service/task_billing.go index 9e43ee2b..29321bd9 100644 --- a/service/task_billing.go +++ b/service/task_billing.go @@ -11,6 +11,7 @@ import ( "github.com/QuantumNous/new-api/model" relaycommon "github.com/QuantumNous/new-api/relay/common" "github.com/QuantumNous/new-api/setting/ratio_setting" + "github.com/QuantumNous/new-api/types" "github.com/gin-gonic/gin" ) @@ -23,9 +24,9 @@ func LogTaskConsumption(c *gin.Context, info *relaycommon.RelayInfo) { if common.StringsContains(constant.TaskPricePatches, info.OriginModelName) { logContent = fmt.Sprintf("%s，按次计费", logContent) } else { - if len(info.PriceData.OtherRatios) > 0 { + if otherRatios := info.PriceData.OtherRatios(); len(otherRatios) > 0 { var contents []string - for key, ra := range info.PriceData.OtherRatios { + for key, ra := range otherRatios { if 1.0 != ra { contents = append(contents, fmt.Sprintf("%s: %.2f", key, ra)) } @@ -126,8 +127,8 @@ func taskBillingOther(task *model.Task) map[string]interface{} { other["model_ratio"] = bc.ModelRatio } other["group_ratio"] = bc.GroupRatio - if len(bc.OtherRatios) > 0 { - for k, v := range bc.OtherRatios { + if priceData := taskBillingContextPriceData(bc); priceData != nil { + for k, v := range priceData.OtherRatios() { other[k] = v } } @@ -140,6 +141,17 @@ func taskBillingOther(task *model.Task) map[string]interface{} { return other } +func taskBillingContextPriceData(bc *model.TaskBillingContext) *types.PriceData { + if bc == nil || len(bc.OtherRatios) == 0 { + return nil + } + priceData := &types.PriceData{} + if !priceData.ReplaceOtherRatios(bc.OtherRatios) { + return nil + } + return priceData +} + // taskModelName 从 BillingContext 或 Properties 中获取模型名称。 func taskModelName(task *model.Task) string { if bc := task.PrivateData.BillingContext; bc != nil && bc.OriginModelName != "" { @@ -294,12 +306,8 @@ func RecalculateTaskQuotaByTokens(ctx context.Context, task *model.Task, totalTo // 计算 OtherRatios 乘积（视频折扣、时长等） otherMultiplier := 1.0 - if bc := task.PrivateData.BillingContext; bc != nil { - for _, r := range bc.OtherRatios { - if r != 1.0 && r > 0 { - otherMultiplier *= r - } - } + if priceData := taskBillingContextPriceData(task.PrivateData.BillingContext); priceData != nil { + otherMultiplier = priceData.OtherRatioMultiplier() } // 计算实际应扣费额度: totalTokens * modelRatio * groupRatio * otherMultiplier（饱和转换，防止溢出成负数） diff --git a/service/task_billing_test.go b/service/task_billing_test.go index b6b1c080..fb1c2915 100644 --- a/service/task_billing_test.go +++ b/service/task_billing_test.go @@ -3,6 +3,7 @@ package service import ( "context" "encoding/json" + "math" "net/http" "os" "testing" @@ -11,7 +12,9 @@ import ( "github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/model" relaycommon "github.com/QuantumNous/new-api/relay/common" + "github.com/QuantumNous/new-api/types" "github.com/glebarez/sqlite" + "github.com/shopspring/decimal" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "gorm.io/gorm" @@ -139,6 +142,102 @@ func makeTask(userId, channelId, quota, tokenId int, billingSource string, subsc } } +func TestPriceDataOtherRatiosFilterAndSnapshot(t *testing.T) { + priceData := types.PriceData{} + + priceData.AddOtherRatio("zero", 0) + priceData.AddOtherRatio("negative", -0.5) + priceData.AddOtherRatio("nan", math.NaN()) + priceData.AddOtherRatio("inf", math.Inf(1)) + priceData.AddOtherRatio("one", 1) + priceData.AddOtherRatio("positive", 2.5) + + ratios := priceData.OtherRatios() + require.Len(t, ratios, 2) + assert.Equal(t, 1.0, ratios["one"]) + assert.Equal(t, 2.5, ratios["positive"]) + assert.True(t, priceData.HasOtherRatio("one")) + assert.False(t, priceData.HasOtherRatio("zero")) + + ratios["positive"] = 99 + ratios["new"] = 3 + nextSnapshot := priceData.OtherRatios() + assert.Equal(t, 2.5, nextSnapshot["positive"]) + assert.NotContains(t, nextSnapshot, "new") +} + +func TestPriceDataReplaceAndApplyOtherRatios(t *testing.T) { + priceData := types.PriceData{} + + replaced := priceData.ReplaceOtherRatios(map[string]float64{ + "zero": 0, + "negative": -3, + "nan": math.NaN(), + "inf": math.Inf(1), + "one": 1, + "duration": 2, + "size": 1.5, + }) + + require.True(t, replaced) + assert.Equal(t, 3.0, priceData.OtherRatioMultiplier()) + assert.Equal(t, 30.0, priceData.ApplyOtherRatiosToFloat(10)) + assert.Equal(t, 10.0, priceData.RemoveOtherRatiosFromFloat(30)) + assert.True(t, decimal.NewFromInt(30).Equal(priceData.ApplyOtherRatiosToDecimal(decimal.NewFromInt(10)))) + + replaced = priceData.ReplaceOtherRatios(map[string]float64{ + "zero": 0, + "nan": math.NaN(), + }) + + require.False(t, replaced) + assert.Nil(t, priceData.OtherRatios()) + assert.Equal(t, 1.0, priceData.OtherRatioMultiplier()) +} + +func TestTaskBillingOtherFiltersHistoricalOtherRatios(t *testing.T) { + task := makeTask(1, 1, 100, 0, BillingSourceWallet, 0) + task.PrivateData.BillingContext.OtherRatios = map[string]float64{ + "seconds": 2, + "identity": 1, + "zero": 0, + "negative": -1, + "nan": math.NaN(), + "inf": math.Inf(1), + } + + other := taskBillingOther(task) + + assert.Equal(t, 2.0, other["seconds"]) + assert.Equal(t, 1.0, other["identity"]) + assert.NotContains(t, other, "zero") + assert.NotContains(t, other, "negative") + assert.NotContains(t, other, "nan") + assert.NotContains(t, other, "inf") +} + +func TestTaskBillingContextPriceDataFiltersMultiplier(t *testing.T) { + priceData := taskBillingContextPriceData(&model.TaskBillingContext{ + OtherRatios: map[string]float64{ + "seconds": 2, + "size": 3, + "identity": 1, + "zero": 0, + "negative": -1, + "nan": math.NaN(), + "inf": math.Inf(1), + }, + }) + + require.NotNil(t, priceData) + assert.Equal(t, 6.0, priceData.OtherRatioMultiplier()) + assert.Equal(t, map[string]float64{ + "seconds": 2, + "size": 3, + "identity": 1, + }, priceData.OtherRatios()) +} + // --------------------------------------------------------------------------- // Read-back helpers // --------------------------------------------------------------------------- diff --git a/service/text_quota.go b/service/text_quota.go index 4f6e4c11..42385ca8 100644 --- a/service/text_quota.go +++ b/service/text_quota.go @@ -296,12 +296,7 @@ func calculateTextQuotaSummary(ctx *gin.Context, relayInfo *relaycommon.RelayInf quotaCalculateDecimal := promptQuota.Add(completionQuota).Mul(ratio) quotaCalculateDecimal = quotaCalculateDecimal.Add(summary.ToolCallSurchargeQuota) quotaCalculateDecimal = quotaCalculateDecimal.Add(audioInputQuota) - - if len(relayInfo.PriceData.OtherRatios) > 0 { - for _, otherRatio := range relayInfo.PriceData.OtherRatios { - quotaCalculateDecimal = quotaCalculateDecimal.Mul(decimal.NewFromFloat(otherRatio)) - } - } + quotaCalculateDecimal = relayInfo.PriceData.ApplyOtherRatiosToDecimal(quotaCalculateDecimal) if !ratio.IsZero() && quotaCalculateDecimal.LessThanOrEqual(decimal.Zero) { quotaCalculateDecimal = decimal.NewFromInt(1) @@ -313,11 +308,7 @@ func calculateTextQuotaSummary(ctx *gin.Context, relayInfo *relaycommon.RelayInf quotaCalculateDecimal := dModelPrice.Mul(dQuotaPerUnit).Mul(dGroupRatio) quotaCalculateDecimal = quotaCalculateDecimal.Add(summary.ToolCallSurchargeQuota) quotaCalculateDecimal = quotaCalculateDecimal.Add(audioInputQuota) - if len(relayInfo.PriceData.OtherRatios) > 0 { - for _, otherRatio := range relayInfo.PriceData.OtherRatios { - quotaCalculateDecimal = quotaCalculateDecimal.Mul(decimal.NewFromFloat(otherRatio)) - } - } + quotaCalculateDecimal = relayInfo.PriceData.ApplyOtherRatiosToDecimal(quotaCalculateDecimal) quota, clamp := common.QuotaFromDecimalChecked(quotaCalculateDecimal) summary.Quota = quota noteQuotaClamp(relayInfo, clamp) diff --git a/types/price_data.go b/types/price_data.go index b4fad00a..8ab1c57c 100644 --- a/types/price_data.go +++ b/types/price_data.go @@ -3,6 +3,8 @@ package types import ( "fmt" "math" + + "github.com/shopspring/decimal" ) type GroupRatioInfo struct { @@ -23,7 +25,7 @@ type PriceData struct { ImageRatio float64 AudioRatio float64 AudioCompletionRatio float64 - OtherRatios map[string]float64 + otherRatios map[string]float64 UsePrice bool Quota int // 按次计费的最终额度（MJ / Task） QuotaToPreConsume int // 按量计费的预消耗额度 @@ -31,15 +33,80 @@ type PriceData struct { } func (p *PriceData) AddOtherRatio(key string, ratio float64) { - if p.OtherRatios == nil { - p.OtherRatios = make(map[string]float64) - } - // NaN/Inf would poison every downstream quota multiplication - // (int(NaN * quota) wraps to a negative charge). - if !(ratio > 0) || math.IsInf(ratio, 1) { + if !isValidOtherRatio(ratio) { return } - p.OtherRatios[key] = ratio + if p.otherRatios == nil { + p.otherRatios = make(map[string]float64) + } + p.otherRatios[key] = ratio +} + +func (p *PriceData) ReplaceOtherRatios(ratios map[string]float64) bool { + p.otherRatios = nil + for key, ratio := range ratios { + p.AddOtherRatio(key, ratio) + } + return len(p.otherRatios) > 0 +} + +func (p *PriceData) HasOtherRatio(key string) bool { + ratio, ok := p.otherRatios[key] + return ok && isValidOtherRatio(ratio) +} + +func (p *PriceData) OtherRatios() map[string]float64 { + if len(p.otherRatios) == 0 { + return nil + } + ratios := make(map[string]float64, len(p.otherRatios)) + for key, ratio := range p.otherRatios { + if isValidOtherRatio(ratio) { + ratios[key] = ratio + } + } + if len(ratios) == 0 { + return nil + } + return ratios +} + +func (p *PriceData) OtherRatioMultiplier() float64 { + multiplier := 1.0 + for _, ratio := range p.otherRatios { + if isValidOtherRatio(ratio) && ratio != 1.0 { + multiplier *= ratio + } + } + return multiplier +} + +func (p *PriceData) ApplyOtherRatiosToFloat(value float64) float64 { + return value * p.OtherRatioMultiplier() +} + +func (p *PriceData) ApplyOtherRatiosToDecimal(value decimal.Decimal) decimal.Decimal { + for _, ratio := range p.otherRatios { + if isValidOtherRatio(ratio) && ratio != 1.0 { + value = value.Mul(decimal.NewFromFloat(ratio)) + } + } + return value +} + +func (p *PriceData) RemoveOtherRatiosFromFloat(value float64) float64 { + for _, ratio := range p.otherRatios { + if isValidOtherRatio(ratio) && ratio != 1.0 { + value /= ratio + } + } + return value +} + +func isValidOtherRatio(ratio float64) bool { + // NaN/Inf would poison every downstream quota multiplication + // (int(NaN * quota) wraps to a negative charge). + return ratio > 0 && !math.IsInf(ratio, 1) } func (p *PriceData) ToSetting() string {