Commit 3fbad6a7 by CaIon

fix(price): add default token estimate for tiered expression pre-consume

parent becc18e3
......@@ -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)
})
}
}
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or sign in to comment