diff --git a/controller/relay.go b/controller/relay.go index 5a321319..7e427086 100644 --- a/controller/relay.go +++ b/controller/relay.go @@ -199,6 +199,10 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) { newAPIError = channelErr break } + if billingErr := service.PrepareTieredBillingForSelectedGroup(c, relayInfo); billingErr != nil { + newAPIError = billingErr + break + } addUsedChannel(c, channel.Id) bodyStorage, bodyErr := common.GetBodyStorage(c) diff --git a/pkg/billingexpr/expr.md b/pkg/billingexpr/expr.md index 89894ab0..c52f2e61 100644 --- a/pkg/billingexpr/expr.md +++ b/pkg/billingexpr/expr.md @@ -159,7 +159,7 @@ When a request arrives and the model uses `tiered_expr` billing: 2. Builds `RequestInput` (headers + body) for `param()` / `header()` functions 3. Runs expression with estimated tokens: `RunExprWithRequest(expr, {P, C}, requestInput)` 4. Converts output to quota: `rawCost / 1,000,000 * QuotaPerUnit` -5. Creates `BillingSnapshot` (frozen state for settlement) and stores on `RelayInfo` +5. Creates `BillingSnapshot` and stores it on `RelayInfo`. Expression and request state stay frozen for settlement. An auto-group retry refreshes group-dependent fields from the selected group before the next upstream attempt. If a free initial group skipped pre-consume and the retry selects a paid group, the billing session is created before that attempt. If an existing session moves to a more expensive group, its reservation is raised to that group's estimate before sending; cheaper groups are refunded only after actual usage is settled. ### 4. Settlement (Actual Billing) @@ -173,7 +173,7 @@ After the upstream response returns with actual token usage: - For Claude-format APIs (input_tokens is text-only): no adjustment needed 2. `TryTieredSettle(relayInfo, params)`: - - Uses the frozen `BillingSnapshot` from pre-consume + - Uses the captured `BillingSnapshot`, whose group-dependent fields have been refreshed from the final selected group - Re-runs the expression with actual token counts - Converts via `quotaConversion()` (version-dispatched) - Returns actual quota diff --git a/pkg/billingexpr/types.go b/pkg/billingexpr/types.go index fa30b5c0..e95399b4 100644 --- a/pkg/billingexpr/types.go +++ b/pkg/billingexpr/types.go @@ -36,8 +36,10 @@ type TraceResult struct { Cost float64 `json:"cost"` } -// BillingSnapshot captures the billing rule state frozen at pre-consume time. -// It is fully serializable and contains no compiled program pointers. +// BillingSnapshot captures billing state at pre-consume time. Expression and +// request fields stay frozen; group-dependent fields are refreshed before an +// auto-group retry and settlement. It is fully serializable and contains no +// compiled program pointers. type BillingSnapshot struct { BillingMode string `json:"billing_mode"` ModelName string `json:"model_name"` diff --git a/relay/common/relay_info.go b/relay/common/relay_info.go index d8d45f99..45ae30bf 100644 --- a/relay/common/relay_info.go +++ b/relay/common/relay_info.go @@ -122,7 +122,7 @@ type RelayInfo struct { // 必须在提交前锁定全额。 ForcePreConsume bool // Billing 是计费会话，封装了预扣费/结算/退款的统一生命周期。 - // 免费模型时为 nil。 + // 初始免费组可为 nil；若 auto 重试切换到付费组，会在发送前创建。 Billing BillingSettler // BillingSource indicates whether this request is billed from wallet quota or subscription. // "" or "wallet" => wallet; "subscription" => subscription @@ -163,8 +163,9 @@ type RelayInfo struct { // It is surfaced onto the consume/task log's admin_info for auditing. QuotaClamp *common.QuotaClamp - // TieredBillingSnapshot is a frozen snapshot of tiered billing rules - // captured at pre-consume time. Non-nil only when billing mode is "tiered_expr". + // TieredBillingSnapshot captures tiered billing rules at pre-consume time. + // Auto-group retries refresh its group-dependent fields before each attempt + // and again before settlement. Non-nil only when billing mode is "tiered_expr". TieredBillingSnapshot *billingexpr.BillingSnapshot BillingRequestInput *billingexpr.RequestInput diff --git a/service/tiered_settle.go b/service/tiered_settle.go index 785ee934..b31d7964 100644 --- a/service/tiered_settle.go +++ b/service/tiered_settle.go @@ -1,9 +1,13 @@ package service import ( + "net/http" + "github.com/QuantumNous/new-api/pkg/billingexpr" relaycommon "github.com/QuantumNous/new-api/relay/common" "github.com/QuantumNous/new-api/relaykit/dto" + "github.com/QuantumNous/new-api/relaykit/types" + "github.com/gin-gonic/gin" ) // TieredResultWrapper wraps billingexpr.TieredResult for use at the service layer. @@ -90,8 +94,60 @@ func BuildTieredTokenParams(usage *dto.Usage, isClaudeUsageSemantic bool, usedVa } } +func refreshTieredBillingGroup(relayInfo *relaycommon.RelayInfo) (*billingexpr.BillingSnapshot, error) { + if relayInfo == nil { + return nil, nil + } + snap := relayInfo.TieredBillingSnapshot + if snap == nil || snap.BillingMode != "tiered_expr" { + return nil, nil + } + + groupRatio := relayInfo.PriceData.GroupRatioInfo.GroupRatio + if snap.GroupRatio == groupRatio { + return snap, nil + } + + estimatedQuotaAfterGroup := snap.EstimatedQuotaBeforeGroup * groupRatio + estimatedQuota, err := billingexpr.QuotaRoundStrict(estimatedQuotaAfterGroup) + if err != nil { + return nil, err + } + snap.GroupRatio = groupRatio + snap.EstimatedQuotaAfterGroup = estimatedQuota + return snap, nil +} + +// PrepareTieredBillingForSelectedGroup refreshes routing-dependent billing +// state before an upstream attempt. An existing session reserves any higher +// estimate before sending. If the initial group was free and skipped +// pre-consume, switching to a paid group creates the session at that point. +func PrepareTieredBillingForSelectedGroup(c *gin.Context, relayInfo *relaycommon.RelayInfo) *types.NewAPIError { + snap, err := refreshTieredBillingGroup(relayInfo) + if err != nil { + return types.NewErrorWithStatusCode( + err, + types.ErrorCodeModelPriceError, + http.StatusBadRequest, + types.ErrOptionWithSkipRetry(), + ) + } + if snap == nil || snap.GroupRatio == 0 { + return nil + } + + if relayInfo.Billing == nil { + return PreConsumeBilling(c, snap.EstimatedQuotaAfterGroup, relayInfo) + } + if err := relayInfo.Billing.Reserve(snap.EstimatedQuotaAfterGroup); err != nil { + return types.NewError(err, types.ErrorCodeUpdateDataError, types.ErrOptionWithSkipRetry()) + } + relayInfo.FinalPreConsumedQuota = relayInfo.Billing.GetPreConsumedQuota() + return nil +} + // TryTieredSettle checks if the request uses tiered_expr billing and, if so, -// computes the actual quota using the frozen BillingSnapshot. Returns: +// computes the actual quota using the captured BillingSnapshot. Returns: // - ok=true, quota, result when tiered billing applies // - ok=false, 0, nil when it doesn't (caller should fall through to existing logic) func TryTieredSettle(relayInfo *relaycommon.RelayInfo, params billingexpr.TokenParams) (ok bool, quota int, result *billingexpr.TieredResult) { diff --git a/service/tiered_settle_test.go b/service/tiered_settle_test.go index de8abd90..fe3a30a0 100644 --- a/service/tiered_settle_test.go +++ b/service/tiered_settle_test.go @@ -5,10 +5,15 @@ import ( "math/rand" "testing" + "github.com/QuantumNous/new-api/model" "github.com/QuantumNous/new-api/pkg/billingexpr" relaycommon "github.com/QuantumNous/new-api/relay/common" "github.com/QuantumNous/new-api/relaykit/dto" + "github.com/QuantumNous/new-api/types" + "github.com/gin-gonic/gin" "github.com/shopspring/decimal" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" ) // Claude Sonnet-style tiered expression: standard vs long-context @@ -309,6 +314,139 @@ func TestTryTieredSettle_NoRequestInput_FallsBackToDefault(t *testing.T) { // Group ratio tests // --------------------------------------------------------------------------- +type recordingBillingSettler struct { + preConsumedQuota int + reserveTargets []int +} + +func (*recordingBillingSettler) Settle(int) error { return nil } + +func (*recordingBillingSettler) Refund(*gin.Context) {} + +func (*recordingBillingSettler) NeedsRefund() bool { return false } + +func (s *recordingBillingSettler) GetPreConsumedQuota() int { + return s.preConsumedQuota +} + +func (s *recordingBillingSettler) Reserve(targetQuota int) error { + s.reserveTargets = append(s.reserveTargets, targetQuota) + if targetQuota > s.preConsumedQuota { + s.preConsumedQuota = targetQuota + } + return nil +} + +func TestPrepareTieredBillingForSelectedGroupUpdatesReservation(t *testing.T) { + const expr = `tier("base", p)` + billing := &recordingBillingSettler{preConsumedQuota: 50_000} + relayInfo := &relaycommon.RelayInfo{ + Billing: billing, + FinalPreConsumedQuota: 50_000, + TieredBillingSnapshot: &billingexpr.BillingSnapshot{ + BillingMode: "tiered_expr", + ExprString: expr, + ExprHash: billingexpr.ExprHashString(expr), + GroupRatio: 0.10, + EstimatedQuotaBeforeGroup: 500_000, + EstimatedQuotaAfterGroup: 50_000, + QuotaPerUnit: testQuotaPerUnit, + }, + PriceData: types.PriceData{ + GroupRatioInfo: types.GroupRatioInfo{GroupRatio: 0.20}, + }, + } + + require.Nil(t, PrepareTieredBillingForSelectedGroup(nil, relayInfo)) + require.Equal(t, []int{100_000}, billing.reserveTargets) + assert.Equal(t, 100_000, billing.preConsumedQuota) + assert.Equal(t, 100_000, relayInfo.FinalPreConsumedQuota) + assert.Equal(t, 0.20, relayInfo.TieredBillingSnapshot.GroupRatio) + assert.Equal(t, 100_000, relayInfo.TieredBillingSnapshot.EstimatedQuotaAfterGroup) +} + +func TestPrepareTieredBillingForSelectedGroupStartsBillingAfterFreeGroup(t *testing.T) { + truncate(t) + gin.SetMode(gin.TestMode) + + const userID = 700 + seedUser(t, userID, 500_000) + + relayInfo := &relaycommon.RelayInfo{ + UserId: userID, + IsPlayground: true, + ForcePreConsume: true, + OriginModelName: "gpt-test", + UserSetting: dto.UserSetting{ + BillingPreference: "wallet_only", + }, + TieredBillingSnapshot: &billingexpr.BillingSnapshot{ + BillingMode: "tiered_expr", + ExprString: `tier("base", p)`, + ExprHash: billingexpr.ExprHashString(`tier("base", p)`), + GroupRatio: 0, + EstimatedQuotaBeforeGroup: 500_000, + QuotaPerUnit: testQuotaPerUnit, + }, + PriceData: types.PriceData{ + GroupRatioInfo: types.GroupRatioInfo{GroupRatio: 0.20}, + }, + } + ctx, _ := gin.CreateTestContext(nil) + + require.Nil(t, PrepareTieredBillingForSelectedGroup(ctx, relayInfo)) + require.NotNil(t, relayInfo.Billing) + assert.Equal(t, 100_000, relayInfo.FinalPreConsumedQuota) + assert.Equal(t, 0.20, relayInfo.TieredBillingSnapshot.GroupRatio) + assert.Equal(t, 100_000, relayInfo.TieredBillingSnapshot.EstimatedQuotaAfterGroup) + + userQuota, err := model.GetUserQuota(userID, false) + require.NoError(t, err) + assert.Equal(t, 400_000, userQuota) +} + +func TestTryTieredSettleUsesFinalGroupAfterRetry(t *testing.T) { + const expr = `tier("base", p)` + tests := []struct { + name string + finalGroupRatio float64 + wantQuota int + }{ + {name: "more expensive final group", finalGroupRatio: 0.20, wantQuota: 100_000}, + {name: "free final group", finalGroupRatio: 0, wantQuota: 0}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + relayInfo := &relaycommon.RelayInfo{ + Billing: &recordingBillingSettler{preConsumedQuota: 50_000}, + FinalPreConsumedQuota: 50_000, + TieredBillingSnapshot: &billingexpr.BillingSnapshot{ + BillingMode: "tiered_expr", + ExprString: expr, + ExprHash: billingexpr.ExprHashString(expr), + GroupRatio: 0.10, + EstimatedQuotaBeforeGroup: 500_000, + EstimatedQuotaAfterGroup: 50_000, + QuotaPerUnit: testQuotaPerUnit, + }, + PriceData: types.PriceData{ + GroupRatioInfo: types.GroupRatioInfo{GroupRatio: tt.finalGroupRatio}, + }, + } + + require.Nil(t, PrepareTieredBillingForSelectedGroup(nil, relayInfo)) + ok, quota, result := TryTieredSettle(relayInfo, billingexpr.TokenParams{P: 1_000_000}) + + require.True(t, ok) + require.NotNil(t, result) + assert.Equal(t, tt.wantQuota, quota) + assert.Equal(t, tt.finalGroupRatio, relayInfo.TieredBillingSnapshot.GroupRatio) + assert.Equal(t, tt.wantQuota, relayInfo.TieredBillingSnapshot.EstimatedQuotaAfterGroup) + }) + } +} + func TestTryTieredSettle_GroupRatioScaling(t *testing.T) { info := makeRelayInfo(flatExpr, 1.5, 1000, 500)