88 lines
3.7 KiB
Go
88 lines
3.7 KiB
Go
package service
|
|
|
|
import (
|
|
"errors"
|
|
"fmt"
|
|
"math"
|
|
|
|
"github.com/QuantumNous/new-api/common"
|
|
"github.com/QuantumNous/new-api/model"
|
|
"github.com/QuantumNous/new-api/pkg/billingexpr"
|
|
benchmark "github.com/QuantumNous/new-api/pkg/model_benchmark"
|
|
fingerprint "github.com/QuantumNous/new-api/pkg/model_fingerprint"
|
|
"github.com/QuantumNous/new-api/relaykit/dto"
|
|
"github.com/QuantumNous/new-api/setting/billing_setting"
|
|
)
|
|
|
|
type DetectionPrice struct {
|
|
Snapshot billingexpr.BillingSnapshot `json:"snapshot"`
|
|
Request billingexpr.RequestInput `json:"request"`
|
|
}
|
|
|
|
func DetectionPricing(userGroup, group, kind string) (DetectionPrice, int, error) {
|
|
expression, ok := billing_setting.GetBillingExprForGroup(group, fingerprint.Model)
|
|
if !ok {
|
|
snapshot, err := model.GetModelPricingSnapshot([]string{fingerprint.Model})
|
|
if err != nil || snapshot == nil || len(snapshot.Entries) != 1 {
|
|
return DetectionPrice{}, 0, errors.New("pricing unavailable")
|
|
}
|
|
converted, err := model.PreviewModelPricingConversion(fingerprint.Model, snapshot.Entries[0].Configured)
|
|
if err != nil || converted == nil || converted.Expression == "" {
|
|
return DetectionPrice{}, 0, errors.New("unsupported detection pricing")
|
|
}
|
|
expression = converted.Expression
|
|
}
|
|
ratio := GetUserGroupRatio(userGroup, group)
|
|
if ratio < 0 || math.IsNaN(ratio) || math.IsInf(ratio, 0) {
|
|
return DetectionPrice{}, 0, errors.New("invalid group ratio")
|
|
}
|
|
price := DetectionPrice{Snapshot: billingexpr.BillingSnapshot{BillingMode: "tiered_expr", ModelName: fingerprint.Model, ExprString: expression, ExprHash: billingexpr.ExprHashString(expression), GroupRatio: ratio, QuotaPerUnit: common.QuotaPerUnit, ExprVersion: 1}}
|
|
effort := "low"
|
|
output, calls := fingerprint.OutputTokens, fingerprint.ProbeCount
|
|
if kind == "benchmark" {
|
|
output, calls = benchmark.OutputTokens, benchmark.ProbeCount
|
|
} else if kind == "vision" {
|
|
output, calls = VisionOutputTokens, 1
|
|
effort = "medium"
|
|
} else if kind != "fingerprint" {
|
|
return price, 0, errors.New("invalid detection type")
|
|
}
|
|
price.Request.Body = []byte(fmt.Sprintf(`{"model":"gpt-6-astra","reasoning":{"effort":"%s"},"max_output_tokens":%d}`, effort, output))
|
|
cost, err := billingexpr.DetectionCostBound(expression, 32768, float64(output))
|
|
if err != nil {
|
|
return price, 0, err
|
|
}
|
|
upper, upperClamp := common.QuotaRoundChecked(math.Ceil(cost / 1_000_000 * common.QuotaPerUnit * ratio))
|
|
if upperClamp != nil {
|
|
return price, 0, errors.New("reservation exceeds quota limit")
|
|
}
|
|
reserve, clamp := common.QuotaFromFloatChecked(float64(upper) * float64(calls))
|
|
if clamp != nil {
|
|
return price, 0, errors.New("reservation exceeds quota limit")
|
|
}
|
|
return price, reserve, nil
|
|
}
|
|
|
|
func DetectionUsageQuota(priceJSON string, input, output, cached int64, cacheWrites ...int64) (int, error) {
|
|
written := int64(0)
|
|
if len(cacheWrites) > 0 {
|
|
written = cacheWrites[0]
|
|
}
|
|
if input < 0 || output < 0 || cached < 0 || written < 0 || cached > input || written > input-cached || input > 32768 || output > 16384 {
|
|
return 0, errors.New("invalid upstream usage")
|
|
}
|
|
var price DetectionPrice
|
|
if err := common.UnmarshalJsonStr(priceJSON, &price); err != nil {
|
|
return 0, err
|
|
}
|
|
usage := &dto.Usage{PromptTokens: int(input), CompletionTokens: int(output)}
|
|
usage.PromptTokensDetails.CachedTokens = int(cached)
|
|
usage.PromptTokensDetails.CachedCreationTokens = int(written)
|
|
params := BuildTieredTokenParams(usage, false, billingexpr.UsedVarsByHash(price.Snapshot.ExprString, price.Snapshot.ExprHash))
|
|
result, err := billingexpr.ComputeTieredQuotaWithRequest(&price.Snapshot, params, price.Request)
|
|
if err != nil || result.Clamp != nil || result.ActualQuotaAfterGroup < 0 {
|
|
return 0, errors.New("usage settlement failed")
|
|
}
|
|
return result.ActualQuotaAfterGroup, nil
|
|
}
|