91ed4e196a
- Added support for tiered billing expressions in the billing system. - Introduced new types and functions for handling billing expressions, including caching and execution. - Updated existing billing logic to accommodate tiered billing scenarios. - Enhanced request handling to support incoming billing expression requests. - Added tests for tiered billing functionality to ensure correctness.
37 lines
1.2 KiB
Go
37 lines
1.2 KiB
Go
package service
|
|
|
|
import (
|
|
"github.com/QuantumNous/new-api/pkg/billingexpr"
|
|
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
|
)
|
|
|
|
// TieredResultWrapper wraps billingexpr.TieredResult for use at the service layer.
|
|
type TieredResultWrapper = billingexpr.TieredResult
|
|
|
|
// TryTieredSettle checks if the request uses tiered_expr billing and, if so,
|
|
// computes the actual quota using the frozen 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) {
|
|
snap := relayInfo.TieredBillingSnapshot
|
|
if snap == nil || snap.BillingMode != "tiered_expr" {
|
|
return false, 0, nil
|
|
}
|
|
|
|
requestInput := billingexpr.RequestInput{}
|
|
if relayInfo.BillingRequestInput != nil {
|
|
requestInput = *relayInfo.BillingRequestInput
|
|
}
|
|
|
|
tr, err := billingexpr.ComputeTieredQuotaWithRequest(snap, params, requestInput)
|
|
if err != nil {
|
|
quota = relayInfo.FinalPreConsumedQuota
|
|
if quota <= 0 {
|
|
quota = snap.EstimatedQuotaAfterGroup
|
|
}
|
|
return true, quota, nil
|
|
}
|
|
|
|
return true, tr.ActualQuotaAfterGroup, &tr
|
|
}
|