Jaymr/new-api
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(billing): improve quota handling and error reporting for pre-consume operations

Browse Source
This commit is contained in:
CaIon
2026-07-11 13:14:22 +08:00
parent 621927f710
commit d9595831bf
12 changed files with 246 additions and 61 deletions
Download Patch File Download Diff File Expand all files Collapse all files
relay/helper/openai_image_request_test.go
+11
View File
@@ -108,6 +108,16 @@ func TestGetAndValidOpenAIImageRequestNBounds(t *testing.T) {
body: fmt.Sprintf(`{"model":"gpt-image-1","prompt":"a cat","n":%d}`, dto.MaxImageN),
wantN: dto.MaxImageN,
},
{
name: "explicit n is accepted",
body: `{"model":"gpt-image-1","prompt":"a cat","n":3}`,
wantN: 3,
},
{
name: "zero n defaults to 1",
body: `{"model":"gpt-image-1","prompt":"a cat","n":0}`,
wantN: 1,
},
{
name: "absent n defaults to 1",
body: `{"model":"gpt-image-1","prompt":"a cat"}`,
@@ -127,6 +137,7 @@ func TestGetAndValidOpenAIImageRequestNBounds(t *testing.T) {
require.NoError(t, err)
require.NotNil(t, req.N)
require.Equal(t, tt.wantN, *req.N)
require.Equal(t, float64(tt.wantN), req.GetTokenCountMeta().BillingRatios["n"])
})
}
relay/helper/price.go
+26 -24
View File
@@ -32,10 +32,6 @@ func modelPriceNotConfiguredError(modelName string, userId int) error {
)
}
func preConsumeQuotaRangeError(modelName string, clamp *common.QuotaClamp) error {
return fmt.Errorf("model %s pre-consume quota is out of range: operation=%s kind=%s value=%g", modelName, clamp.Op, clamp.Kind, clamp.Original)
}
// https://docs.claude.com/en/docs/build-with-claude/prompt-caching#1-hour-cache-duration
const claudeCacheCreation1hMultiplier = 6 / 3.75
@@ -121,20 +117,15 @@ func ModelPriceHelper(c *gin.Context, info *relaycommon.RelayInfo, promptTokens
audioRatio = ratio_setting.GetAudioRatio(info.OriginModelName)
audioCompletionRatio = ratio_setting.GetAudioCompletionRatio(info.OriginModelName)
ratio := modelRatio * groupRatioInfo.GroupRatio
var clamp *common.QuotaClamp
preConsumedQuota, clamp = common.QuotaFromFloatChecked(float64(preConsumedTokens) * ratio)
if clamp != nil {
return types.PriceData{}, preConsumeQuotaRangeError(info.OriginModelName, clamp)
quota, err := common.QuotaFromFloatStrict(float64(preConsumedTokens) * ratio)
if err != nil {
return types.PriceData{}, err
}
preConsumedQuota = quota
} else {
if meta.ImagePriceRatio != 0 {
modelPrice = modelPrice * meta.ImagePriceRatio
}
var clamp *common.QuotaClamp
preConsumedQuota, clamp = common.QuotaFromFloatChecked(modelPrice * common.QuotaPerUnit * groupRatioInfo.GroupRatio)
if clamp != nil {
return types.PriceData{}, preConsumeQuotaRangeError(info.OriginModelName, clamp)
}
}
// check if free model pre-consume is disabled
@@ -172,6 +163,17 @@ func ModelPriceHelper(c *gin.Context, info *relaycommon.RelayInfo, promptTokens
CacheCreation1hRatio: cacheCreationRatio1h,
QuotaToPreConsume: preConsumedQuota,
}
if usePrice {
for name, ratio := range meta.BillingRatios {
priceData.AddOtherRatio(name, ratio)
}
quotaToPreConsume := priceData.ApplyOtherRatiosToFloat(modelPrice * common.QuotaPerUnit * groupRatioInfo.GroupRatio)
quota, err := common.QuotaFromFloatStrict(quotaToPreConsume)
if err != nil {
return types.PriceData{}, err
}
priceData.QuotaToPreConsume = quota
}
if common.DebugEnabled {
logger.LogDebug(c, "model_price_helper result: %s", priceData.ToSetting())
@@ -211,10 +213,10 @@ func ModelPriceHelperPerCall(c *gin.Context, info *relaycommon.RelayInfo) (types
freeModel := false
if usePrice {
var clamp *common.QuotaClamp
quota, clamp = common.QuotaFromFloatChecked(modelPrice * common.QuotaPerUnit * groupRatioInfo.GroupRatio)
if clamp != nil {
return types.PriceData{}, preConsumeQuotaRangeError(info.OriginModelName, clamp)
var err error
quota, err = common.QuotaFromFloatStrict(modelPrice * common.QuotaPerUnit * groupRatioInfo.GroupRatio)
if err != nil {
return types.PriceData{}, err
}
if !operation_setting.GetQuotaSetting().EnableFreeModelPreConsume {
if groupRatioInfo.GroupRatio == 0 || modelPrice == 0 {
@@ -224,10 +226,10 @@ func ModelPriceHelperPerCall(c *gin.Context, info *relaycommon.RelayInfo) (types
}
} else {
// 按量计费：以模型倍率的一半作为预扣额度
var clamp *common.QuotaClamp
quota, clamp = common.QuotaFromFloatChecked(modelRatio / 2 * common.QuotaPerUnit * groupRatioInfo.GroupRatio)
if clamp != nil {
return types.PriceData{}, preConsumeQuotaRangeError(info.OriginModelName, clamp)
var err error
quota, err = common.QuotaFromFloatStrict(modelRatio / 2 * common.QuotaPerUnit * groupRatioInfo.GroupRatio)
if err != nil {
return types.PriceData{}, err
}
modelPrice = -1
if !operation_setting.GetQuotaSetting().EnableFreeModelPreConsume {
@@ -290,9 +292,9 @@ func modelPriceHelperTiered(c *gin.Context, info *relaycommon.RelayInfo, promptT
// Expression coefficients are $/1M tokens prices; convert to quota the same way per-call billing does.
quotaBeforeGroup := rawCost / 1_000_000 * common.QuotaPerUnit
preConsumedQuota, clamp := billingexpr.QuotaRoundChecked(quotaBeforeGroup * groupRatioInfo.GroupRatio)
if clamp != nil {
return types.PriceData{}, preConsumeQuotaRangeError(info.OriginModelName, clamp)
preConsumedQuota, err := billingexpr.QuotaRoundStrict(quotaBeforeGroup * groupRatioInfo.GroupRatio)
if err != nil {
return types.PriceData{}, err
}
freeModel := false
relay/helper/price_test.go
+100 -2
View File
@@ -10,6 +10,7 @@ import (
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/setting/billing_setting"
"github.com/QuantumNous/new-api/setting/config"
"github.com/QuantumNous/new-api/setting/ratio_setting"
"github.com/QuantumNous/new-api/types"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
@@ -52,7 +53,9 @@ func TestModelPriceHelperTieredUsesPreloadedRequestInput(t *testing.T) {
},
}
priceData, err := ModelPriceHelper(ctx, info, 1000, &types.TokenCountMeta{})
priceData, err := ModelPriceHelper(ctx, info, 1000, &types.TokenCountMeta{
BillingRatios: map[string]float64{"n": 3},
})
require.NoError(t, err)
require.Equal(t, 1500, priceData.QuotaToPreConsume)
require.NotNil(t, info.TieredBillingSnapshot)
@@ -172,5 +175,100 @@ func TestModelPriceHelperTieredRejectsPreConsumeOverflow(t *testing.T) {
_, err := ModelPriceHelper(ctx, info, 1000, &types.TokenCountMeta{})
require.ErrorContains(t, err, "pre-consume quota is out of range")
var clamp *common.QuotaClamp
require.ErrorAs(t, err, &clamp)
require.Equal(t, "QuotaRound", clamp.Op)
require.Equal(t, common.QuotaClampOverflow, clamp.Kind)
}
func TestModelPriceHelperRequestBillingRatiosOnlyApplyToFixedPrice(t *testing.T) {
gin.SetMode(gin.TestMode)
savedModelPrices := ratio_setting.ModelPrice2JSONString()
savedModelRatios := ratio_setting.ModelRatio2JSONString()
t.Cleanup(func() {
require.NoError(t, ratio_setting.UpdateModelPriceByJSONString(savedModelPrices))
require.NoError(t, ratio_setting.UpdateModelRatioByJSONString(savedModelRatios))
})
modelPrices, err := common.Marshal(map[string]float64{
"fixed-image-price": 0.04,
"fractional-image-price": 0.0000012,
"overflow-image-price": float64(common.MaxQuota) / common.QuotaPerUnit / 2,
})
require.NoError(t, err)
require.NoError(t, ratio_setting.UpdateModelPriceByJSONString(string(modelPrices)))
modelRatios, err := common.Marshal(map[string]float64{"ratio-image-price": 15})
require.NoError(t, err)
require.NoError(t, ratio_setting.UpdateModelRatioByJSONString(string(modelRatios)))
tests := []struct {
name string
model string
wantQuota int
wantUsePrice bool
wantImageCount bool
}{
{
name: "fixed price applies image count",
model: "fixed-image-price",
wantQuota: 180000,
wantUsePrice: true,
wantImageCount: true,
},
{
name: "ratio price ignores request billing ratios",
model: "ratio-image-price",
wantQuota: 15000,
wantUsePrice: false,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
ctx, _ := gin.CreateTestContext(httptest.NewRecorder())
ctx.Set("group", "default")
info := &relaycommon.RelayInfo{
OriginModelName: tt.model,
UserGroup: "default",
UsingGroup: "default",
}
meta := &types.TokenCountMeta{
ImagePriceRatio: 3,
BillingRatios: map[string]float64{"n": 3},
}
priceData, err := ModelPriceHelper(ctx, info, 1000, meta)
require.NoError(t, err)
require.Equal(t, tt.wantQuota, priceData.QuotaToPreConsume)
require.Equal(t, tt.wantUsePrice, priceData.UsePrice)
require.Equal(t, tt.wantImageCount, priceData.HasOtherRatio("n"))
require.Equal(t, priceData.OtherRatios(), info.PriceData.OtherRatios())
})
}
newInfo := func(model string) (*gin.Context, *relaycommon.RelayInfo) {
ctx, _ := gin.CreateTestContext(httptest.NewRecorder())
ctx.Set("group", "default")
return ctx, &relaycommon.RelayInfo{
OriginModelName: model,
UserGroup: "default",
UsingGroup: "default",
}
}
meta := &types.TokenCountMeta{BillingRatios: map[string]float64{"n": 3}}
ctx, info := newInfo("fractional-image-price")
priceData, err := ModelPriceHelper(ctx, info, 0, meta)
require.NoError(t, err)
// 0.0000012 * 500000 * 3 = 1.8, then truncate once to 1.
require.Equal(t, 1, priceData.QuotaToPreConsume)
ctx, info = newInfo("overflow-image-price")
_, err = ModelPriceHelper(ctx, info, 0, meta)
var clamp *common.QuotaClamp
require.ErrorAs(t, err, &clamp)
require.Equal(t, "QuotaFromFloat", clamp.Op)
require.Equal(t, common.QuotaClampOverflow, clamp.Kind)
require.Nil(t, info.Billing)
}
relay/image_handler.go
-10
View File
@@ -123,16 +123,6 @@ func ImageHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *type
imageN = *request.N
}
// n is handled via OtherRatio so it is applied exactly once in quota
// calculation (both price-based and ratio-based paths).
// 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 !info.PriceData.HasOtherRatio("n") {
info.PriceData.AddOtherRatio("n", float64(imageN))
}
}
if usage.(*dto.Usage).TotalTokens == 0 {
usage.(*dto.Usage).TotalTokens = 1
}