mirror of
https://github.com/QuantumNous/new-api.git
synced 2026-09-12 15:21:09 +00:00
fix(price): add default token estimate for tiered expression pre-consume
This commit is contained in:
@@ -35,6 +35,11 @@ func modelPriceNotConfiguredError(modelName string, userId int) error {
|
|||||||
// 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
|
||||||
|
|
||||||
|
// defaultTieredPreConsumeMaxTokens is the fallback completion-token estimate
|
||||||
|
// used for tiered expression pre-consume when the client omits max_tokens, so
|
||||||
|
// the pre-consumed quota still reflects a plausible output cost in paid groups.
|
||||||
|
const defaultTieredPreConsumeMaxTokens = 8192
|
||||||
|
|
||||||
// HandleGroupRatio checks for "auto_group" in the context and updates the group ratio and relayInfo.UsingGroup if present
|
// HandleGroupRatio checks for "auto_group" in the context and updates the group ratio and relayInfo.UsingGroup if present
|
||||||
func HandleGroupRatio(ctx *gin.Context, relayInfo *relaycommon.RelayInfo) types.GroupRatioInfo {
|
func HandleGroupRatio(ctx *gin.Context, relayInfo *relaycommon.RelayInfo) types.GroupRatioInfo {
|
||||||
groupRatioInfo := types.GroupRatioInfo{
|
groupRatioInfo := types.GroupRatioInfo{
|
||||||
@@ -244,9 +249,9 @@ func modelPriceHelperTiered(c *gin.Context, info *relaycommon.RelayInfo, promptT
|
|||||||
return types.PriceData{}, fmt.Errorf("model %s is configured as tiered_expr but has no billing expression", info.OriginModelName)
|
return types.PriceData{}, fmt.Errorf("model %s is configured as tiered_expr but has no billing expression", info.OriginModelName)
|
||||||
}
|
}
|
||||||
|
|
||||||
estimatedCompletionTokens := 0
|
estimatedCompletionTokens := meta.MaxTokens
|
||||||
if meta.MaxTokens != 0 {
|
if estimatedCompletionTokens == 0 && groupRatioInfo.GroupRatio != 0 {
|
||||||
estimatedCompletionTokens = meta.MaxTokens
|
estimatedCompletionTokens = defaultTieredPreConsumeMaxTokens
|
||||||
}
|
}
|
||||||
|
|
||||||
requestInput, err := ResolveIncomingBillingExprRequestInput(c, info)
|
requestInput, err := ResolveIncomingBillingExprRequestInput(c, info)
|
||||||
|
|||||||
@@ -60,3 +60,81 @@ func TestModelPriceHelperTieredUsesPreloadedRequestInput(t *testing.T) {
|
|||||||
require.Equal(t, billing_setting.BillingModeTieredExpr, info.TieredBillingSnapshot.BillingMode)
|
require.Equal(t, billing_setting.BillingModeTieredExpr, info.TieredBillingSnapshot.BillingMode)
|
||||||
require.Equal(t, common.QuotaPerUnit, info.TieredBillingSnapshot.QuotaPerUnit)
|
require.Equal(t, common.QuotaPerUnit, info.TieredBillingSnapshot.QuotaPerUnit)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func TestModelPriceHelperTieredPreConsumeMaxTokensFallback(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-fallback-model":"tiered_expr"}`,
|
||||||
|
"billing_setting.billing_expr": `{"tiered-fallback-model":"tier(\"base\", p * 3 + c * 15)"}`,
|
||||||
|
"group_ratio_setting.group_ratio": `{"default":1,"free":0}`,
|
||||||
|
}))
|
||||||
|
|
||||||
|
const promptTokens = 1000
|
||||||
|
|
||||||
|
cases := []struct {
|
||||||
|
name string
|
||||||
|
group string
|
||||||
|
maxTokens int
|
||||||
|
expected int
|
||||||
|
}{
|
||||||
|
{
|
||||||
|
// max_tokens omitted in a paid group -> fall back to 8192 completion tokens.
|
||||||
|
// p*3 + c*15 = 1000*3 + 8192*15 = 125880 -> /1e6 * 500000 = 62940
|
||||||
|
name: "non-free group falls back to 8192 completion tokens",
|
||||||
|
group: "default",
|
||||||
|
maxTokens: 0,
|
||||||
|
expected: 62940,
|
||||||
|
},
|
||||||
|
{
|
||||||
|
// explicit max_tokens is used verbatim, no fallback.
|
||||||
|
// 1000*3 + 100*15 = 4500 -> /1e6 * 500000 = 2250
|
||||||
|
name: "explicit max_tokens is used verbatim",
|
||||||
|
group: "default",
|
||||||
|
maxTokens: 100,
|
||||||
|
expected: 2250,
|
||||||
|
},
|
||||||
|
{
|
||||||
|
// free group (ratio 0) stays zero; fallback is gated on non-zero group ratio.
|
||||||
|
name: "free group stays zero without fallback",
|
||||||
|
group: "free",
|
||||||
|
maxTokens: 0,
|
||||||
|
expected: 0,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, tc := range cases {
|
||||||
|
t.Run(tc.name, func(t *testing.T) {
|
||||||
|
recorder := httptest.NewRecorder()
|
||||||
|
ctx, _ := gin.CreateTestContext(recorder)
|
||||||
|
req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
|
||||||
|
req.Header.Set("Content-Type", "application/json")
|
||||||
|
ctx.Request = req
|
||||||
|
ctx.Set("group", tc.group)
|
||||||
|
|
||||||
|
info := &relaycommon.RelayInfo{
|
||||||
|
OriginModelName: "tiered-fallback-model",
|
||||||
|
UserGroup: tc.group,
|
||||||
|
UsingGroup: tc.group,
|
||||||
|
RequestHeaders: map[string]string{"Content-Type": "application/json"},
|
||||||
|
BillingRequestInput: &billingexpr.RequestInput{
|
||||||
|
Headers: map[string]string{"Content-Type": "application/json"},
|
||||||
|
Body: []byte(`{}`),
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
priceData, err := ModelPriceHelper(ctx, info, promptTokens, &types.TokenCountMeta{MaxTokens: tc.maxTokens})
|
||||||
|
require.NoError(t, err)
|
||||||
|
require.Equal(t, tc.expected, priceData.QuotaToPreConsume)
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user