cneicy
ceddfd3d6b
Sync production source including routing, billing, Astra detection and public uptime. Update custom build instructions and license packaging.
288 lines
11 KiB
Go
288 lines
11 KiB
Go
package model
|
|
|
|
import (
|
|
"errors"
|
|
"fmt"
|
|
"time"
|
|
|
|
"github.com/QuantumNous/new-api/common"
|
|
"gorm.io/gorm"
|
|
"gorm.io/gorm/clause"
|
|
)
|
|
|
|
// RelayAttempt is the durable financial fact and log outbox for one upstream
|
|
// generation. It contains numeric usage and pricing, never a prompt or key.
|
|
type RelayAttempt struct {
|
|
ID int64 `json:"id"`
|
|
AttemptKey string `json:"attempt_key" gorm:"type:varchar(96);uniqueIndex"`
|
|
RequestID string `json:"request_id" gorm:"type:varchar(64);index"`
|
|
Number int `json:"number"`
|
|
UserID int `json:"-" gorm:"index"`
|
|
TokenID int `json:"-"`
|
|
ChannelID int `json:"channel_id"`
|
|
GroupName string `json:"group"`
|
|
RequestedModel string `json:"requested_model"`
|
|
ReturnedModel string `json:"returned_model"`
|
|
GuardResult string `json:"guard_result"`
|
|
Funding string `json:"billing_source"`
|
|
SubscriptionID int `json:"-"`
|
|
SubscriptionResetTime int64 `json:"-"`
|
|
Reserved int `json:"-"`
|
|
Quota int `json:"quota"`
|
|
ObservedQuota int `json:"-"`
|
|
EstimateQuota int `json:"-"`
|
|
Billable bool `json:"-"`
|
|
UsageSource string `json:"usage_source"`
|
|
State string `json:"state" gorm:"type:varchar(24);index"`
|
|
UsageJSON string `json:"-" gorm:"type:text"`
|
|
PriceJSON string `json:"-" gorm:"type:text"`
|
|
LogJSON string `json:"-" gorm:"type:text"`
|
|
LogDelivered bool `json:"-" gorm:"index"`
|
|
CreatedAt int64 `json:"created_at" gorm:"autoCreateTime:false"`
|
|
UpdatedAt int64 `json:"-" gorm:"autoUpdateTime:false"`
|
|
Deadline int64 `json:"-" gorm:"index"`
|
|
}
|
|
|
|
// SQLite cannot upgrade a read transaction while another attempt is writing.
|
|
// Retry the entire rolled-back transaction, never an individual debit statement.
|
|
func relayAttemptTransaction(fn func(*gorm.DB) error) error {
|
|
for retry := 0; ; retry++ {
|
|
err := DB.Transaction(fn)
|
|
var coded interface{ Code() int }
|
|
busy := errors.As(err, &coded) && (coded.Code()&255 == 5 || coded.Code()&255 == 6)
|
|
if !busy || retry >= 24 {
|
|
return err
|
|
}
|
|
delay := time.Duration(retry+1) * 5 * time.Millisecond
|
|
if delay > 100*time.Millisecond {
|
|
delay = 100 * time.Millisecond
|
|
}
|
|
time.Sleep(delay)
|
|
}
|
|
}
|
|
|
|
var ErrAttemptQuota = errors.New("insufficient quota for another upstream attempt")
|
|
|
|
func ReserveRelayAttempt(attempt *RelayAttempt, preference string) error {
|
|
if attempt.Reserved < 0 || attempt.Reserved > common.MaxWalletQuota {
|
|
return ErrAttemptQuota
|
|
}
|
|
attempt.CreatedAt = time.Now().Unix()
|
|
attempt.UpdatedAt = attempt.CreatedAt
|
|
attempt.State = "reserved"
|
|
var tokenKey string
|
|
original := *attempt
|
|
err := relayAttemptTransaction(func(tx *gorm.DB) error {
|
|
*attempt = original
|
|
var existing RelayAttempt
|
|
if err := tx.Where("attempt_key = ?", attempt.AttemptKey).First(&existing).Error; err == nil {
|
|
*attempt = existing
|
|
return nil
|
|
} else if !errors.Is(err, gorm.ErrRecordNotFound) {
|
|
return err
|
|
}
|
|
var user User
|
|
if err := lockForUpdate(tx).First(&user, attempt.UserID).Error; err != nil {
|
|
return err
|
|
}
|
|
var token Token
|
|
if err := lockForUpdate(tx).First(&token, attempt.TokenID).Error; err != nil {
|
|
return err
|
|
}
|
|
tokenKey = token.Key
|
|
if token.UserId != attempt.UserID || token.Status != common.TokenStatusEnabled || (!token.UnlimitedQuota && token.RemainQuota < attempt.Reserved) {
|
|
return ErrAttemptQuota
|
|
}
|
|
if err := invalidateTokenCacheForMutation(token.Key); err != nil {
|
|
return err
|
|
}
|
|
amount := int64(attempt.Reserved)
|
|
wallet := func() bool { return user.Quota >= attempt.Reserved }
|
|
trySub := func() (bool, bool, error) {
|
|
var subs []UserSubscription
|
|
if err := lockForUpdate(tx).Where("user_id = ? AND status = ? AND end_time > ?", attempt.UserID, "active", attempt.CreatedAt).Order("end_time asc, id asc").Find(&subs).Error; err != nil {
|
|
return false, false, err
|
|
}
|
|
allow := len(subs) == 0
|
|
for _, sub := range subs {
|
|
allow = allow || sub.AllowWalletOverflow
|
|
plan, err := getSubscriptionPlanByIdTx(tx, sub.PlanId)
|
|
if err != nil {
|
|
return false, allow, err
|
|
}
|
|
if err := maybeResetUserSubscriptionWithPlanTx(tx, &sub, plan, attempt.CreatedAt); err != nil {
|
|
return false, allow, err
|
|
}
|
|
if sub.AmountTotal > 0 && sub.AmountTotal-sub.AmountUsed < amount {
|
|
continue
|
|
}
|
|
if err := tx.Model(&UserSubscription{}).Where("id = ?", sub.Id).Update("amount_used", gorm.Expr("amount_used + ?", amount)).Error; err != nil {
|
|
return false, allow, err
|
|
}
|
|
attempt.Funding = "subscription"
|
|
attempt.SubscriptionID = sub.Id
|
|
attempt.SubscriptionResetTime = sub.LastResetTime
|
|
return true, allow, nil
|
|
}
|
|
return false, allow, nil
|
|
}
|
|
useWallet := false
|
|
switch common.NormalizeBillingPreference(preference) {
|
|
case "wallet_only":
|
|
useWallet = wallet()
|
|
case "wallet_first":
|
|
useWallet = wallet()
|
|
if !useWallet {
|
|
ok, _, err := trySub()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if !ok {
|
|
return ErrAttemptQuota
|
|
}
|
|
}
|
|
case "subscription_only":
|
|
ok, _, err := trySub()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if !ok {
|
|
return ErrAttemptQuota
|
|
}
|
|
default:
|
|
ok, allow, err := trySub()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if !ok {
|
|
useWallet = allow && wallet()
|
|
}
|
|
}
|
|
if useWallet {
|
|
attempt.Funding = "wallet"
|
|
if err := tx.Model(&User{}).Where("id = ?", user.Id).Update("quota", gorm.Expr("quota - ?", attempt.Reserved)).Error; err != nil {
|
|
return err
|
|
}
|
|
} else if attempt.Funding != "subscription" {
|
|
return ErrAttemptQuota
|
|
}
|
|
if err := tx.Model(&Token{}).Where("id = ?", token.Id).Updates(map[string]any{"remain_quota": gorm.Expr("remain_quota - ?", attempt.Reserved), "used_quota": gorm.Expr("used_quota + ?", attempt.Reserved), "accessed_time": attempt.CreatedAt}).Error; err != nil {
|
|
return err
|
|
}
|
|
return tx.Create(attempt).Error
|
|
})
|
|
if err == nil {
|
|
_ = invalidateUserCache(attempt.UserID)
|
|
_ = invalidateTokenCacheForMutation(tokenKey)
|
|
}
|
|
return err
|
|
}
|
|
|
|
func CheckpointRelayAttempt(id int64, quota int, usage, logJSON, returned, guard, source string, billable bool, deadline int64) error {
|
|
return DB.Model(&RelayAttempt{}).Where("id = ? AND state <> ?", id, "settled").Updates(map[string]any{"observed_quota": quota, "usage_json": usage, "log_json": logJSON, "returned_model": returned, "guard_result": guard, "usage_source": source, "billable": billable, "state": "collecting", "updated_at": time.Now().Unix(), "deadline": deadline}).Error
|
|
}
|
|
|
|
// Settlement and all balances/statistics commit together. Retrying after any
|
|
// process crash cannot repeat a debit, including the delta-zero case.
|
|
func SettleRelayAttempt(id int64, quota int, log *Log, details ...RelayAttempt) error {
|
|
if quota < 0 || quota > common.MaxWalletQuota {
|
|
return fmt.Errorf("invalid attempt charge")
|
|
}
|
|
var attempt RelayAttempt
|
|
var tokenKey string
|
|
err := relayAttemptTransaction(func(tx *gorm.DB) error {
|
|
attempt = RelayAttempt{}
|
|
if err := lockForUpdate(tx).First(&attempt, id).Error; err != nil {
|
|
return err
|
|
}
|
|
if attempt.State == "settled" {
|
|
return nil
|
|
}
|
|
delta := quota - attempt.Reserved
|
|
var user User
|
|
if err := lockForUpdate(tx).Unscoped().First(&user, attempt.UserID).Error; err != nil {
|
|
return err
|
|
}
|
|
var token Token
|
|
if err := lockForUpdate(tx).Unscoped().First(&token, attempt.TokenID).Error; err != nil {
|
|
return err
|
|
}
|
|
tokenKey = token.Key
|
|
if err := invalidateTokenCacheForMutation(token.Key); err != nil {
|
|
return err
|
|
}
|
|
if attempt.Funding == "subscription" {
|
|
if err := tx.Model(&UserSubscription{}).Where("id = ? AND last_reset_time = ?", attempt.SubscriptionID, attempt.SubscriptionResetTime).Update("amount_used", gorm.Expr("CASE WHEN amount_used + ? < 0 THEN 0 ELSE amount_used + ? END", delta, delta)).Error; err != nil {
|
|
return err
|
|
}
|
|
} else if err := tx.Unscoped().Model(&User{}).Where("id = ?", attempt.UserID).Update("quota", gorm.Expr("quota - ?", delta)).Error; err != nil {
|
|
return err
|
|
}
|
|
if err := tx.Unscoped().Model(&Token{}).Where("id = ?", attempt.TokenID).Updates(map[string]any{"remain_quota": gorm.Expr("remain_quota - ?", delta), "used_quota": gorm.Expr("used_quota + ?", delta)}).Error; err != nil {
|
|
return err
|
|
}
|
|
if quota > 0 || attempt.Billable || (len(details) > 0 && details[0].Billable) {
|
|
if err := tx.Unscoped().Model(&User{}).Where("id = ?", attempt.UserID).Updates(map[string]any{"used_quota": gorm.Expr("used_quota + ?", quota), "request_count": gorm.Expr("request_count + 1")}).Error; err != nil {
|
|
return err
|
|
}
|
|
if err := tx.Model(&Channel{}).Where("id = ?", attempt.ChannelID).Update("used_quota", gorm.Expr("used_quota + ?", quota)).Error; err != nil {
|
|
return err
|
|
}
|
|
}
|
|
if len(details) > 0 {
|
|
d := details[0]
|
|
attempt.UsageJSON = d.UsageJSON
|
|
attempt.ReturnedModel = d.ReturnedModel
|
|
attempt.GuardResult = d.GuardResult
|
|
attempt.UsageSource = d.UsageSource
|
|
attempt.Billable = d.Billable
|
|
attempt.ObservedQuota = quota
|
|
}
|
|
attempt.Quota = quota
|
|
attempt.State = "settled"
|
|
attempt.UpdatedAt = time.Now().Unix()
|
|
if log != nil {
|
|
log.Quota = quota
|
|
log.AttemptKey = &attempt.AttemptKey
|
|
if common.DataExportEnabled {
|
|
if err := tx.Create(&QuotaData{AttemptKey: &attempt.AttemptKey, UserID: attempt.UserID, Username: log.Username, ModelName: log.ModelName, CreatedAt: log.CreatedAt - (log.CreatedAt % 3600), UseGroup: attempt.GroupName, TokenID: attempt.TokenID, ChannelID: attempt.ChannelID, NodeName: common.NodeName, Count: 1, Quota: quota, TokenUsed: log.PromptTokens + log.CompletionTokens}).Error; err != nil {
|
|
return err
|
|
}
|
|
}
|
|
body, err := common.Marshal(log)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
attempt.LogJSON = string(body)
|
|
}
|
|
return tx.Save(&attempt).Error
|
|
})
|
|
if err == nil {
|
|
_ = invalidateUserCache(attempt.UserID)
|
|
_ = invalidateTokenCacheForMutation(tokenKey)
|
|
}
|
|
return err
|
|
}
|
|
|
|
func DeliverRelayAttemptLogs() error {
|
|
var attempts []RelayAttempt
|
|
if err := DB.Where("state = ? AND log_delivered = ? AND log_json <> ?", "settled", false, "").Limit(100).Find(&attempts).Error; err != nil {
|
|
return err
|
|
}
|
|
for _, a := range attempts {
|
|
var log Log
|
|
if err := common.UnmarshalJsonStr(a.LogJSON, &log); err != nil {
|
|
return err
|
|
}
|
|
log.Id = 0
|
|
log.AttemptKey = &a.AttemptKey
|
|
if err := LOG_DB.Clauses(clause.OnConflict{Columns: []clause.Column{{Name: "attempt_key"}}, DoNothing: true}).Create(&log).Error; err != nil {
|
|
return err
|
|
}
|
|
if err := DB.Model(&RelayAttempt{}).Where("id = ?", a.ID).Update("log_delivered", true).Error; err != nil {
|
|
return err
|
|
}
|
|
}
|
|
return nil
|
|
}
|