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
|
||||
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
|
||||
func HandleGroupRatio(ctx *gin.Context, relayInfo *relaycommon.RelayInfo) 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)
|
||||
}
|
||||
|
||||
estimatedCompletionTokens := 0
|
||||
if meta.MaxTokens != 0 {
|
||||
estimatedCompletionTokens = meta.MaxTokens
|
||||
estimatedCompletionTokens := meta.MaxTokens
|
||||
if estimatedCompletionTokens == 0 && groupRatioInfo.GroupRatio != 0 {
|
||||
estimatedCompletionTokens = defaultTieredPreConsumeMaxTokens
|
||||
}
|
||||
|
||||
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, 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