mirror of
https://github.com/QuantumNous/new-api.git
synced 2026-09-13 15:54:34 +00:00
fix(billing): reject saturated pre-consume quota
This commit is contained in:
@@ -12,3 +12,10 @@ import "github.com/QuantumNous/new-api/common"
|
|||||||
func QuotaRound(f float64) int {
|
func QuotaRound(f float64) int {
|
||||||
return common.QuotaRound(f)
|
return common.QuotaRound(f)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// QuotaRoundChecked is QuotaRound but also reports whether the result had to
|
||||||
|
// be saturated. Pre-consume callers use this to reject an unrepresentable
|
||||||
|
// estimate before any quota is deducted.
|
||||||
|
func QuotaRoundChecked(f float64) (int, *common.QuotaClamp) {
|
||||||
|
return common.QuotaRoundChecked(f)
|
||||||
|
}
|
||||||
|
|||||||
+28
-5
@@ -32,6 +32,10 @@ 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
|
// https://docs.claude.com/en/docs/build-with-claude/prompt-caching#1-hour-cache-duration
|
||||||
const claudeCacheCreation1hMultiplier = 6 / 3.75
|
const claudeCacheCreation1hMultiplier = 6 / 3.75
|
||||||
|
|
||||||
@@ -117,12 +121,20 @@ func ModelPriceHelper(c *gin.Context, info *relaycommon.RelayInfo, promptTokens
|
|||||||
audioRatio = ratio_setting.GetAudioRatio(info.OriginModelName)
|
audioRatio = ratio_setting.GetAudioRatio(info.OriginModelName)
|
||||||
audioCompletionRatio = ratio_setting.GetAudioCompletionRatio(info.OriginModelName)
|
audioCompletionRatio = ratio_setting.GetAudioCompletionRatio(info.OriginModelName)
|
||||||
ratio := modelRatio * groupRatioInfo.GroupRatio
|
ratio := modelRatio * groupRatioInfo.GroupRatio
|
||||||
preConsumedQuota = common.QuotaFromFloat(float64(preConsumedTokens) * ratio)
|
var clamp *common.QuotaClamp
|
||||||
|
preConsumedQuota, clamp = common.QuotaFromFloatChecked(float64(preConsumedTokens) * ratio)
|
||||||
|
if clamp != nil {
|
||||||
|
return types.PriceData{}, preConsumeQuotaRangeError(info.OriginModelName, clamp)
|
||||||
|
}
|
||||||
} else {
|
} else {
|
||||||
if meta.ImagePriceRatio != 0 {
|
if meta.ImagePriceRatio != 0 {
|
||||||
modelPrice = modelPrice * meta.ImagePriceRatio
|
modelPrice = modelPrice * meta.ImagePriceRatio
|
||||||
}
|
}
|
||||||
preConsumedQuota = common.QuotaFromFloat(modelPrice * common.QuotaPerUnit * groupRatioInfo.GroupRatio)
|
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
|
// check if free model pre-consume is disabled
|
||||||
@@ -199,7 +211,11 @@ func ModelPriceHelperPerCall(c *gin.Context, info *relaycommon.RelayInfo) (types
|
|||||||
freeModel := false
|
freeModel := false
|
||||||
|
|
||||||
if usePrice {
|
if usePrice {
|
||||||
quota = common.QuotaFromFloat(modelPrice * common.QuotaPerUnit * groupRatioInfo.GroupRatio)
|
var clamp *common.QuotaClamp
|
||||||
|
quota, clamp = common.QuotaFromFloatChecked(modelPrice * common.QuotaPerUnit * groupRatioInfo.GroupRatio)
|
||||||
|
if clamp != nil {
|
||||||
|
return types.PriceData{}, preConsumeQuotaRangeError(info.OriginModelName, clamp)
|
||||||
|
}
|
||||||
if !operation_setting.GetQuotaSetting().EnableFreeModelPreConsume {
|
if !operation_setting.GetQuotaSetting().EnableFreeModelPreConsume {
|
||||||
if groupRatioInfo.GroupRatio == 0 || modelPrice == 0 {
|
if groupRatioInfo.GroupRatio == 0 || modelPrice == 0 {
|
||||||
quota = 0
|
quota = 0
|
||||||
@@ -208,7 +224,11 @@ func ModelPriceHelperPerCall(c *gin.Context, info *relaycommon.RelayInfo) (types
|
|||||||
}
|
}
|
||||||
} else {
|
} else {
|
||||||
// 按量计费:以模型倍率的一半作为预扣额度
|
// 按量计费:以模型倍率的一半作为预扣额度
|
||||||
quota = common.QuotaFromFloat(modelRatio / 2 * common.QuotaPerUnit * groupRatioInfo.GroupRatio)
|
var clamp *common.QuotaClamp
|
||||||
|
quota, clamp = common.QuotaFromFloatChecked(modelRatio / 2 * common.QuotaPerUnit * groupRatioInfo.GroupRatio)
|
||||||
|
if clamp != nil {
|
||||||
|
return types.PriceData{}, preConsumeQuotaRangeError(info.OriginModelName, clamp)
|
||||||
|
}
|
||||||
modelPrice = -1
|
modelPrice = -1
|
||||||
if !operation_setting.GetQuotaSetting().EnableFreeModelPreConsume {
|
if !operation_setting.GetQuotaSetting().EnableFreeModelPreConsume {
|
||||||
if groupRatioInfo.GroupRatio == 0 || modelRatio == 0 {
|
if groupRatioInfo.GroupRatio == 0 || modelRatio == 0 {
|
||||||
@@ -270,7 +290,10 @@ 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.
|
// Expression coefficients are $/1M tokens prices; convert to quota the same way per-call billing does.
|
||||||
quotaBeforeGroup := rawCost / 1_000_000 * common.QuotaPerUnit
|
quotaBeforeGroup := rawCost / 1_000_000 * common.QuotaPerUnit
|
||||||
preConsumedQuota := billingexpr.QuotaRound(quotaBeforeGroup * groupRatioInfo.GroupRatio)
|
preConsumedQuota, clamp := billingexpr.QuotaRoundChecked(quotaBeforeGroup * groupRatioInfo.GroupRatio)
|
||||||
|
if clamp != nil {
|
||||||
|
return types.PriceData{}, preConsumeQuotaRangeError(info.OriginModelName, clamp)
|
||||||
|
}
|
||||||
|
|
||||||
freeModel := false
|
freeModel := false
|
||||||
if !operation_setting.GetQuotaSetting().EnableFreeModelPreConsume {
|
if !operation_setting.GetQuotaSetting().EnableFreeModelPreConsume {
|
||||||
|
|||||||
@@ -138,3 +138,39 @@ func TestModelPriceHelperTieredPreConsumeMaxTokensFallback(t *testing.T) {
|
|||||||
})
|
})
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func TestModelPriceHelperTieredRejectsPreConsumeOverflow(t *testing.T) {
|
||||||
|
gin.SetMode(gin.TestMode)
|
||||||
|
|
||||||
|
saved := map[string]string{}
|
||||||
|
require.NoError(t, config.GlobalConfig.SaveToDB(func(key, value string) error {
|
||||||
|
saved[key] = value
|
||||||
|
return nil
|
||||||
|
}))
|
||||||
|
t.Cleanup(func() {
|
||||||
|
require.NoError(t, config.GlobalConfig.LoadFromDB(saved))
|
||||||
|
})
|
||||||
|
|
||||||
|
require.NoError(t, config.GlobalConfig.LoadFromDB(map[string]string{
|
||||||
|
"billing_setting.billing_mode": `{"tiered-overflow-model":"tiered_expr"}`,
|
||||||
|
"billing_setting.billing_expr": `{"tiered-overflow-model":"tier(\"overflow\", p * 1000000000000000)"}`,
|
||||||
|
"group_ratio_setting.group_ratio": `{"default":1}`,
|
||||||
|
}))
|
||||||
|
|
||||||
|
recorder := httptest.NewRecorder()
|
||||||
|
ctx, _ := gin.CreateTestContext(recorder)
|
||||||
|
ctx.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
|
||||||
|
ctx.Set("group", "default")
|
||||||
|
info := &relaycommon.RelayInfo{
|
||||||
|
OriginModelName: "tiered-overflow-model",
|
||||||
|
UserGroup: "default",
|
||||||
|
UsingGroup: "default",
|
||||||
|
BillingRequestInput: &billingexpr.RequestInput{
|
||||||
|
Body: []byte(`{}`),
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
_, err := ModelPriceHelper(ctx, info, 1000, &types.TokenCountMeta{})
|
||||||
|
|
||||||
|
require.ErrorContains(t, err, "pre-consume quota is out of range")
|
||||||
|
}
|
||||||
|
|||||||
@@ -2,6 +2,7 @@ package service
|
|||||||
|
|
||||||
import (
|
import (
|
||||||
"fmt"
|
"fmt"
|
||||||
|
"net/http"
|
||||||
|
|
||||||
"github.com/QuantumNous/new-api/logger"
|
"github.com/QuantumNous/new-api/logger"
|
||||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||||
@@ -17,6 +18,23 @@ const (
|
|||||||
// PreConsumeBilling 根据用户计费偏好创建 BillingSession 并执行预扣费。
|
// PreConsumeBilling 根据用户计费偏好创建 BillingSession 并执行预扣费。
|
||||||
// 会话存储在 relayInfo.Billing 上,供后续 Settle / Refund 使用。
|
// 会话存储在 relayInfo.Billing 上,供后续 Settle / Refund 使用。
|
||||||
func PreConsumeBilling(c *gin.Context, preConsumedQuota int, relayInfo *relaycommon.RelayInfo) *types.NewAPIError {
|
func PreConsumeBilling(c *gin.Context, preConsumedQuota int, relayInfo *relaycommon.RelayInfo) *types.NewAPIError {
|
||||||
|
if relayInfo != nil && relayInfo.QuotaClamp != nil {
|
||||||
|
clamp := relayInfo.QuotaClamp
|
||||||
|
return types.NewErrorWithStatusCode(
|
||||||
|
fmt.Errorf("pre-consume quota is out of range: operation=%s kind=%s value=%g", clamp.Op, clamp.Kind, clamp.Original),
|
||||||
|
types.ErrorCodeModelPriceError,
|
||||||
|
http.StatusBadRequest,
|
||||||
|
types.ErrOptionWithSkipRetry(),
|
||||||
|
)
|
||||||
|
}
|
||||||
|
if preConsumedQuota < 0 {
|
||||||
|
return types.NewErrorWithStatusCode(
|
||||||
|
fmt.Errorf("pre-consume quota cannot be negative: %d", preConsumedQuota),
|
||||||
|
types.ErrorCodeModelPriceError,
|
||||||
|
http.StatusBadRequest,
|
||||||
|
types.ErrOptionWithSkipRetry(),
|
||||||
|
)
|
||||||
|
}
|
||||||
session, apiErr := NewBillingSession(c, relayInfo, preConsumedQuota)
|
session, apiErr := NewBillingSession(c, relayInfo, preConsumedQuota)
|
||||||
if apiErr != nil {
|
if apiErr != nil {
|
||||||
return apiErr
|
return apiErr
|
||||||
|
|||||||
@@ -1,10 +1,12 @@
|
|||||||
package service
|
package service
|
||||||
|
|
||||||
import (
|
import (
|
||||||
|
"net/http"
|
||||||
"testing"
|
"testing"
|
||||||
|
|
||||||
"github.com/QuantumNous/new-api/common"
|
"github.com/QuantumNous/new-api/common"
|
||||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||||
|
"github.com/QuantumNous/new-api/types"
|
||||||
|
|
||||||
"github.com/gin-gonic/gin"
|
"github.com/gin-gonic/gin"
|
||||||
"github.com/stretchr/testify/require"
|
"github.com/stretchr/testify/require"
|
||||||
@@ -72,3 +74,36 @@ func TestAttachQuotaSaturationNoClampNoMarker(t *testing.T) {
|
|||||||
_, hasAdmin := other["admin_info"]
|
_, hasAdmin := other["admin_info"]
|
||||||
require.False(t, hasAdmin, "no admin_info should be added when there is no clamp")
|
require.False(t, hasAdmin, "no admin_info should be added when there is no clamp")
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func TestPreConsumeBillingRejectsSaturatedQuotaBeforeDeduction(t *testing.T) {
|
||||||
|
gin.SetMode(gin.TestMode)
|
||||||
|
c, _ := gin.CreateTestContext(nil)
|
||||||
|
info := &relaycommon.RelayInfo{
|
||||||
|
QuotaClamp: &common.QuotaClamp{
|
||||||
|
Op: "QuotaFromFloat",
|
||||||
|
Kind: common.QuotaClampOverflow,
|
||||||
|
Original: 1e30,
|
||||||
|
Clamped: common.MaxQuota,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
apiErr := PreConsumeBilling(c, common.MaxQuota, info)
|
||||||
|
|
||||||
|
require.NotNil(t, apiErr)
|
||||||
|
require.Equal(t, types.ErrorCodeModelPriceError, apiErr.GetErrorCode())
|
||||||
|
require.Equal(t, http.StatusBadRequest, apiErr.StatusCode)
|
||||||
|
require.Nil(t, info.Billing)
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestPreConsumeBillingRejectsNegativeQuotaBeforeDeduction(t *testing.T) {
|
||||||
|
gin.SetMode(gin.TestMode)
|
||||||
|
c, _ := gin.CreateTestContext(nil)
|
||||||
|
info := &relaycommon.RelayInfo{}
|
||||||
|
|
||||||
|
apiErr := PreConsumeBilling(c, -1, info)
|
||||||
|
|
||||||
|
require.NotNil(t, apiErr)
|
||||||
|
require.Equal(t, types.ErrorCodeModelPriceError, apiErr.GetErrorCode())
|
||||||
|
require.Equal(t, http.StatusBadRequest, apiErr.StatusCode)
|
||||||
|
require.Nil(t, info.Billing)
|
||||||
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user