Ecosystem/ai-gateway
5
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: make quota int64

Browse Source
This commit is contained in:
JustSong 2024-03-13 20:00:51 +08:00
parent a72e5fcc9e
commit e99150bdb9
16 changed files with 70 additions and 70 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
common/config/config.go
View File

@ -76,14 +76,14 @@ var MessagePusherToken = ""
var TurnstileSiteKey = "" var TurnstileSiteKey = ""
var TurnstileSecretKey = "" var TurnstileSecretKey = ""
var QuotaForNewUser = 0 var QuotaForNewUser int64 = 0
var QuotaForInviter = 0 var QuotaForInviter int64 = 0
var QuotaForInvitee = 0 var QuotaForInvitee int64 = 0
var ChannelDisableThreshold = 5.0 var ChannelDisableThreshold = 5.0
var AutomaticDisableChannelEnabled = false var AutomaticDisableChannelEnabled = false
var AutomaticEnableChannelEnabled = false var AutomaticEnableChannelEnabled = false
var QuotaRemindThreshold = 1000 var QuotaRemindThreshold int64 = 1000
var PreConsumedQuota = 500 var PreConsumedQuota int64 = 500
var ApproximateTokenEnabled = false var ApproximateTokenEnabled = false
var RetryTimes = 0 var RetryTimes = 0

2
common/utils.go
View File

@ -5,7 +5,7 @@ import (
"github.com/songquanpeng/one-api/common/config" "github.com/songquanpeng/one-api/common/config"
) )
func LogQuota(quota int) string { func LogQuota(quota int64) string {
if config.DisplayInCurrencyEnabled { if config.DisplayInCurrencyEnabled {
return fmt.Sprintf("＄%.6f 额度", float64(quota)/config.QuotaPerUnit) return fmt.Sprintf("＄%.6f 额度", float64(quota)/config.QuotaPerUnit)
} else { } else {

6
controller/billing.go
View File

@ -8,8 +8,8 @@ import (
) )
func GetSubscription(c *gin.Context) { func GetSubscription(c *gin.Context) {
var remainQuota int var remainQuota int64
var usedQuota int var usedQuota int64
var err error var err error
var token *model.Token var token *model.Token
var expiredTime int64 var expiredTime int64
@ -60,7 +60,7 @@ func GetSubscription(c *gin.Context) {
} }
func GetUsage(c *gin.Context) { func GetUsage(c *gin.Context) {
var quota int var quota int64
var err error var err error
var token *model.Token var token *model.Token
if config.DisplayTokenStatEnabled { if config.DisplayTokenStatEnabled {

8
model/cache.go
View File

@ -71,7 +71,7 @@ func CacheGetUserGroup(id int) (group string, err error) {
return group, err return group, err
} }
func fetchAndUpdateUserQuota(ctx context.Context, id int) (quota int, err error) { func fetchAndUpdateUserQuota(ctx context.Context, id int) (quota int64, err error) {
quota, err = GetUserQuota(id) quota, err = GetUserQuota(id)
if err != nil { if err != nil {
return 0, err return 0, err
@ -83,7 +83,7 @@ func fetchAndUpdateUserQuota(ctx context.Context, id int) (quota int, err error)
return return
} }
func CacheGetUserQuota(ctx context.Context, id int) (quota int, err error) { func CacheGetUserQuota(ctx context.Context, id int) (quota int64, err error) {
if !common.RedisEnabled { if !common.RedisEnabled {
return GetUserQuota(id) return GetUserQuota(id)
} }
@ -91,7 +91,7 @@ func CacheGetUserQuota(ctx context.Context, id int) (quota int, err error) {
if err != nil { if err != nil {
return fetchAndUpdateUserQuota(ctx, id) return fetchAndUpdateUserQuota(ctx, id)
} }
quota, err = strconv.Atoi(quotaString) quota, err = strconv.ParseInt(quotaString, 10, 64)
if err != nil { if err != nil {
return 0, nil return 0, nil
} }
@ -114,7 +114,7 @@ func CacheUpdateUserQuota(ctx context.Context, id int) error {
return err return err
} }
func CacheDecreaseUserQuota(id int, quota int) error { func CacheDecreaseUserQuota(id int, quota int64) error {
if !common.RedisEnabled { if !common.RedisEnabled {
return nil return nil
} }

4
model/channel.go
View File

@ -178,7 +178,7 @@ func UpdateChannelStatusById(id int, status int) {
} }
} }
func UpdateChannelUsedQuota(id int, quota int) { func UpdateChannelUsedQuota(id int, quota int64) {
if config.BatchUpdateEnabled { if config.BatchUpdateEnabled {
addNewRecord(BatchUpdateTypeChannelUsedQuota, id, quota) addNewRecord(BatchUpdateTypeChannelUsedQuota, id, quota)
return return
@ -186,7 +186,7 @@ func UpdateChannelUsedQuota(id int, quota int) {
updateChannelUsedQuota(id, quota) updateChannelUsedQuota(id, quota)
} }
func updateChannelUsedQuota(id int, quota int) { func updateChannelUsedQuota(id int, quota int64) {
err := DB.Model(&Channel{}).Where("id = ?", id).Update("used_quota", gorm.Expr("used_quota + ?", quota)).Error err := DB.Model(&Channel{}).Where("id = ?", id).Update("used_quota", gorm.Expr("used_quota + ?", quota)).Error
if err != nil { if err != nil {
logger.SysError("failed to update channel used quota: " + err.Error()) logger.SysError("failed to update channel used quota: " + err.Error())

6
model/log.go
View File

@ -51,7 +51,7 @@ func RecordLog(userId int, logType int, content string) {
} }
} }
func RecordConsumeLog(ctx context.Context, userId int, channelId int, promptTokens int, completionTokens int, modelName string, tokenName string, quota int, content string) { func RecordConsumeLog(ctx context.Context, userId int, channelId int, promptTokens int, completionTokens int, modelName string, tokenName string, quota int64, content string) {
logger.Info(ctx, fmt.Sprintf("record consume log: userId=%d, channelId=%d, promptTokens=%d, completionTokens=%d, modelName=%s, tokenName=%s, quota=%d, content=%s", userId, channelId, promptTokens, completionTokens, modelName, tokenName, quota, content)) logger.Info(ctx, fmt.Sprintf("record consume log: userId=%d, channelId=%d, promptTokens=%d, completionTokens=%d, modelName=%s, tokenName=%s, quota=%d, content=%s", userId, channelId, promptTokens, completionTokens, modelName, tokenName, quota, content))
if !config.LogConsumeEnabled { if !config.LogConsumeEnabled {
return return
@ -66,7 +66,7 @@ func RecordConsumeLog(ctx context.Context, userId int, channelId int, promptToke
CompletionTokens: completionTokens, CompletionTokens: completionTokens,
TokenName: tokenName, TokenName: tokenName,
ModelName: modelName, ModelName: modelName,
Quota: quota, Quota: int(quota),
ChannelId: channelId, ChannelId: channelId,
} }
err := DB.Create(log).Error err := DB.Create(log).Error
@ -137,7 +137,7 @@ func SearchUserLogs(userId int, keyword string) (logs []*Log, err error) {
return logs, err return logs, err
} }
func SumUsedQuota(logType int, startTimestamp int64, endTimestamp int64, modelName string, username string, tokenName string, channel int) (quota int) { func SumUsedQuota(logType int, startTimestamp int64, endTimestamp int64, modelName string, username string, tokenName string, channel int) (quota int64) {
tx := DB.Table("logs").Select("ifnull(sum(quota),0)") tx := DB.Table("logs").Select("ifnull(sum(quota),0)")
if username != "" { if username != "" {
tx = tx.Where("username = ?", username) tx = tx.Where("username = ?", username)

20
model/option.go
View File

@ -61,11 +61,11 @@ func InitOptionMap() {
config.OptionMap["MessagePusherToken"] = "" config.OptionMap["MessagePusherToken"] = ""
config.OptionMap["TurnstileSiteKey"] = "" config.OptionMap["TurnstileSiteKey"] = ""
config.OptionMap["TurnstileSecretKey"] = "" config.OptionMap["TurnstileSecretKey"] = ""
config.OptionMap["QuotaForNewUser"] = strconv.Itoa(config.QuotaForNewUser) config.OptionMap["QuotaForNewUser"] = strconv.FormatInt(config.QuotaForNewUser, 10)
config.OptionMap["QuotaForInviter"] = strconv.Itoa(config.QuotaForInviter) config.OptionMap["QuotaForInviter"] = strconv.FormatInt(config.QuotaForInviter, 10)
config.OptionMap["QuotaForInvitee"] = strconv.Itoa(config.QuotaForInvitee) config.OptionMap["QuotaForInvitee"] = strconv.FormatInt(config.QuotaForInvitee, 10)
config.OptionMap["QuotaRemindThreshold"] = strconv.Itoa(config.QuotaRemindThreshold) config.OptionMap["QuotaRemindThreshold"] = strconv.FormatInt(config.QuotaRemindThreshold, 10)
config.OptionMap["PreConsumedQuota"] = strconv.Itoa(config.PreConsumedQuota) config.OptionMap["PreConsumedQuota"] = strconv.FormatInt(config.PreConsumedQuota, 10)
config.OptionMap["ModelRatio"] = common.ModelRatio2JSONString() config.OptionMap["ModelRatio"] = common.ModelRatio2JSONString()
config.OptionMap["GroupRatio"] = common.GroupRatio2JSONString() config.OptionMap["GroupRatio"] = common.GroupRatio2JSONString()
config.OptionMap["CompletionRatio"] = common.CompletionRatio2JSONString() config.OptionMap["CompletionRatio"] = common.CompletionRatio2JSONString()
@ -193,15 +193,15 @@ func updateOptionMap(key string, value string) (err error) {
case "TurnstileSecretKey": case "TurnstileSecretKey":
config.TurnstileSecretKey = value config.TurnstileSecretKey = value
case "QuotaForNewUser": case "QuotaForNewUser":
config.QuotaForNewUser, _ = strconv.Atoi(value) config.QuotaForNewUser, _ = strconv.ParseInt(value, 10, 64)
case "QuotaForInviter": case "QuotaForInviter":
config.QuotaForInviter, _ = strconv.Atoi(value) config.QuotaForInviter, _ = strconv.ParseInt(value, 10, 64)
case "QuotaForInvitee": case "QuotaForInvitee":
config.QuotaForInvitee, _ = strconv.Atoi(value) config.QuotaForInvitee, _ = strconv.ParseInt(value, 10, 64)
case "QuotaRemindThreshold": case "QuotaRemindThreshold":
config.QuotaRemindThreshold, _ = strconv.Atoi(value) config.QuotaRemindThreshold, _ = strconv.ParseInt(value, 10, 64)
case "PreConsumedQuota": case "PreConsumedQuota":
config.PreConsumedQuota, _ = strconv.Atoi(value) config.PreConsumedQuota, _ = strconv.ParseInt(value, 10, 64)
case "RetryTimes": case "RetryTimes":
config.RetryTimes, _ = strconv.Atoi(value) config.RetryTimes, _ = strconv.Atoi(value)
case "ModelRatio": case "ModelRatio":

4
model/redemption.go
View File

@ -14,7 +14,7 @@ type Redemption struct {
Key string `json:"key" gorm:"type:char(32);uniqueIndex"` Key string `json:"key" gorm:"type:char(32);uniqueIndex"`
Status int `json:"status" gorm:"default:1"` Status int `json:"status" gorm:"default:1"`
Name string `json:"name" gorm:"index"` Name string `json:"name" gorm:"index"`
Quota int `json:"quota" gorm:"default:100"` Quota int64 `json:"quota" gorm:"default:100"`
CreatedTime int64 `json:"created_time" gorm:"bigint"` CreatedTime int64 `json:"created_time" gorm:"bigint"`
RedeemedTime int64 `json:"redeemed_time" gorm:"bigint"` RedeemedTime int64 `json:"redeemed_time" gorm:"bigint"`
Count int `json:"count" gorm:"-:all"` // only for api request Count int `json:"count" gorm:"-:all"` // only for api request
@ -42,7 +42,7 @@ func GetRedemptionById(id int) (*Redemption, error) {
return &redemption, err return &redemption, err
} }
func Redeem(key string, userId int) (quota int, err error) { func Redeem(key string, userId int) (quota int64, err error) {
if key == "" { if key == "" {
return 0, errors.New("未提供兑换码") return 0, errors.New("未提供兑换码")
} }

16
model/token.go
View File

@ -20,9 +20,9 @@ type Token struct {
CreatedTime int64 `json:"created_time" gorm:"bigint"` CreatedTime int64 `json:"created_time" gorm:"bigint"`
AccessedTime int64 `json:"accessed_time" gorm:"bigint"` AccessedTime int64 `json:"accessed_time" gorm:"bigint"`
ExpiredTime int64 `json:"expired_time" gorm:"bigint;default:-1"` // -1 means never expired ExpiredTime int64 `json:"expired_time" gorm:"bigint;default:-1"` // -1 means never expired
RemainQuota int `json:"remain_quota" gorm:"default:0"` RemainQuota int64 `json:"remain_quota" gorm:"default:0"`
UnlimitedQuota bool `json:"unlimited_quota" gorm:"default:false"` UnlimitedQuota bool `json:"unlimited_quota" gorm:"default:false"`
UsedQuota int `json:"used_quota" gorm:"default:0"` // used quota UsedQuota int64 `json:"used_quota" gorm:"default:0"` // used quota
} }
func GetAllUserTokens(userId int, startIdx int, num int) ([]*Token, error) { func GetAllUserTokens(userId int, startIdx int, num int) ([]*Token, error) {
@ -138,7 +138,7 @@ func DeleteTokenById(id int, userId int) (err error) {
return token.Delete() return token.Delete()
} }
func IncreaseTokenQuota(id int, quota int) (err error) { func IncreaseTokenQuota(id int, quota int64) (err error) {
if quota < 0 { if quota < 0 {
return errors.New("quota 不能为负数！") return errors.New("quota 不能为负数！")
} }
@ -149,7 +149,7 @@ func IncreaseTokenQuota(id int, quota int) (err error) {
return increaseTokenQuota(id, quota) return increaseTokenQuota(id, quota)
} }
func increaseTokenQuota(id int, quota int) (err error) { func increaseTokenQuota(id int, quota int64) (err error) {
err = DB.Model(&Token{}).Where("id = ?", id).Updates( err = DB.Model(&Token{}).Where("id = ?", id).Updates(
map[string]interface{}{ map[string]interface{}{
"remain_quota": gorm.Expr("remain_quota + ?", quota), "remain_quota": gorm.Expr("remain_quota + ?", quota),
@ -160,7 +160,7 @@ func increaseTokenQuota(id int, quota int) (err error) {
return err return err
} }
func DecreaseTokenQuota(id int, quota int) (err error) { func DecreaseTokenQuota(id int, quota int64) (err error) {
if quota < 0 { if quota < 0 {
return errors.New("quota 不能为负数！") return errors.New("quota 不能为负数！")
} }
@ -171,7 +171,7 @@ func DecreaseTokenQuota(id int, quota int) (err error) {
return decreaseTokenQuota(id, quota) return decreaseTokenQuota(id, quota)
} }
func decreaseTokenQuota(id int, quota int) (err error) { func decreaseTokenQuota(id int, quota int64) (err error) {
err = DB.Model(&Token{}).Where("id = ?", id).Updates( err = DB.Model(&Token{}).Where("id = ?", id).Updates(
map[string]interface{}{ map[string]interface{}{
"remain_quota": gorm.Expr("remain_quota - ?", quota), "remain_quota": gorm.Expr("remain_quota - ?", quota),
@ -182,7 +182,7 @@ func decreaseTokenQuota(id int, quota int) (err error) {
return err return err
} }
func PreConsumeTokenQuota(tokenId int, quota int) (err error) { func PreConsumeTokenQuota(tokenId int, quota int64) (err error) {
if quota < 0 { if quota < 0 {
return errors.New("quota 不能为负数！") return errors.New("quota 不能为负数！")
} }
@ -232,7 +232,7 @@ func PreConsumeTokenQuota(tokenId int, quota int) (err error) {
return err return err
} }
func PostConsumeTokenQuota(tokenId int, quota int) (err error) { func PostConsumeTokenQuota(tokenId int, quota int64) (err error) {
token, err := GetTokenById(tokenId) token, err := GetTokenById(tokenId)
if quota > 0 { if quota > 0 {
err = DecreaseUserQuota(token.UserId, quota) err = DecreaseUserQuota(token.UserId, quota)

22
model/user.go
View File

@ -26,8 +26,8 @@ type User struct {
WeChatId string `json:"wechat_id" gorm:"column:wechat_id;index"` WeChatId string `json:"wechat_id" gorm:"column:wechat_id;index"`
VerificationCode string `json:"verification_code" gorm:"-:all"` // this field is only for Email verification, don't save it to database! VerificationCode string `json:"verification_code" gorm:"-:all"` // this field is only for Email verification, don't save it to database!
AccessToken string `json:"access_token" gorm:"type:char(32);column:access_token;uniqueIndex"` // this token is for system management AccessToken string `json:"access_token" gorm:"type:char(32);column:access_token;uniqueIndex"` // this token is for system management
Quota int `json:"quota" gorm:"type:int;default:0"` Quota int64 `json:"quota" gorm:"type:int;default:0"`
UsedQuota int `json:"used_quota" gorm:"type:int;default:0;column:used_quota"` // used quota UsedQuota int64 `json:"used_quota" gorm:"type:int;default:0;column:used_quota"` // used quota
RequestCount int `json:"request_count" gorm:"type:int;default:0;"` // request number RequestCount int `json:"request_count" gorm:"type:int;default:0;"` // request number
Group string `json:"group" gorm:"type:varchar(32);default:'default'"` Group string `json:"group" gorm:"type:varchar(32);default:'default'"`
AffCode string `json:"aff_code" gorm:"type:varchar(32);column:aff_code;uniqueIndex"` AffCode string `json:"aff_code" gorm:"type:varchar(32);column:aff_code;uniqueIndex"`
@ -274,12 +274,12 @@ func ValidateAccessToken(token string) (user *User) {
return nil return nil
} }
func GetUserQuota(id int) (quota int, err error) { func GetUserQuota(id int) (quota int64, err error) {
err = DB.Model(&User{}).Where("id = ?", id).Select("quota").Find(&quota).Error err = DB.Model(&User{}).Where("id = ?", id).Select("quota").Find(&quota).Error
return quota, err return quota, err
} }
func GetUserUsedQuota(id int) (quota int, err error) { func GetUserUsedQuota(id int) (quota int64, err error) {
err = DB.Model(&User{}).Where("id = ?", id).Select("used_quota").Find(&quota).Error err = DB.Model(&User{}).Where("id = ?", id).Select("used_quota").Find(&quota).Error
return quota, err return quota, err
} }
@ -299,7 +299,7 @@ func GetUserGroup(id int) (group string, err error) {
return group, err return group, err
} }
func IncreaseUserQuota(id int, quota int) (err error) { func IncreaseUserQuota(id int, quota int64) (err error) {
if quota < 0 { if quota < 0 {
return errors.New("quota 不能为负数！") return errors.New("quota 不能为负数！")
} }
@ -310,12 +310,12 @@ func IncreaseUserQuota(id int, quota int) (err error) {
return increaseUserQuota(id, quota) return increaseUserQuota(id, quota)
} }
func increaseUserQuota(id int, quota int) (err error) { func increaseUserQuota(id int, quota int64) (err error) {
err = DB.Model(&User{}).Where("id = ?", id).Update("quota", gorm.Expr("quota + ?", quota)).Error err = DB.Model(&User{}).Where("id = ?", id).Update("quota", gorm.Expr("quota + ?", quota)).Error
return err return err
} }
func DecreaseUserQuota(id int, quota int) (err error) { func DecreaseUserQuota(id int, quota int64) (err error) {
if quota < 0 { if quota < 0 {
return errors.New("quota 不能为负数！") return errors.New("quota 不能为负数！")
} }
@ -326,7 +326,7 @@ func DecreaseUserQuota(id int, quota int) (err error) {
return decreaseUserQuota(id, quota) return decreaseUserQuota(id, quota)
} }
func decreaseUserQuota(id int, quota int) (err error) { func decreaseUserQuota(id int, quota int64) (err error) {
err = DB.Model(&User{}).Where("id = ?", id).Update("quota", gorm.Expr("quota - ?", quota)).Error err = DB.Model(&User{}).Where("id = ?", id).Update("quota", gorm.Expr("quota - ?", quota)).Error
return err return err
} }
@ -336,7 +336,7 @@ func GetRootUserEmail() (email string) {
return email return email
} }
func UpdateUserUsedQuotaAndRequestCount(id int, quota int) { func UpdateUserUsedQuotaAndRequestCount(id int, quota int64) {
if config.BatchUpdateEnabled { if config.BatchUpdateEnabled {
addNewRecord(BatchUpdateTypeUsedQuota, id, quota) addNewRecord(BatchUpdateTypeUsedQuota, id, quota)
addNewRecord(BatchUpdateTypeRequestCount, id, 1) addNewRecord(BatchUpdateTypeRequestCount, id, 1)
@ -345,7 +345,7 @@ func UpdateUserUsedQuotaAndRequestCount(id int, quota int) {
updateUserUsedQuotaAndRequestCount(id, quota, 1) updateUserUsedQuotaAndRequestCount(id, quota, 1)
} }
func updateUserUsedQuotaAndRequestCount(id int, quota int, count int) { func updateUserUsedQuotaAndRequestCount(id int, quota int64, count int) {
err := DB.Model(&User{}).Where("id = ?", id).Updates( err := DB.Model(&User{}).Where("id = ?", id).Updates(
map[string]interface{}{ map[string]interface{}{
"used_quota": gorm.Expr("used_quota + ?", quota), "used_quota": gorm.Expr("used_quota + ?", quota),
@ -357,7 +357,7 @@ func updateUserUsedQuotaAndRequestCount(id int, quota int, count int) {
} }
} }
func updateUserUsedQuota(id int, quota int) { func updateUserUsedQuota(id int, quota int64) {
err := DB.Model(&User{}).Where("id = ?", id).Updates( err := DB.Model(&User{}).Where("id = ?", id).Updates(
map[string]interface{}{ map[string]interface{}{
"used_quota": gorm.Expr("used_quota + ?", quota), "used_quota": gorm.Expr("used_quota + ?", quota),

10
model/utils.go
View File

@ -16,12 +16,12 @@ const (
BatchUpdateTypeCount // if you add a new type, you need to add a new map and a new lock BatchUpdateTypeCount // if you add a new type, you need to add a new map and a new lock
) )
var batchUpdateStores []map[int]int var batchUpdateStores []map[int]int64
var batchUpdateLocks []sync.Mutex var batchUpdateLocks []sync.Mutex
func init() { func init() {
for i := 0; i < BatchUpdateTypeCount; i++ { for i := 0; i < BatchUpdateTypeCount; i++ {
batchUpdateStores = append(batchUpdateStores, make(map[int]int)) batchUpdateStores = append(batchUpdateStores, make(map[int]int64))
batchUpdateLocks = append(batchUpdateLocks, sync.Mutex{}) batchUpdateLocks = append(batchUpdateLocks, sync.Mutex{})
} }
} }
@ -35,7 +35,7 @@ func InitBatchUpdater() {
}() }()
} }
func addNewRecord(type_ int, id int, value int) { func addNewRecord(type_ int, id int, value int64) {
batchUpdateLocks[type_].Lock() batchUpdateLocks[type_].Lock()
defer batchUpdateLocks[type_].Unlock() defer batchUpdateLocks[type_].Unlock()
if _, ok := batchUpdateStores[type_][id]; !ok { if _, ok := batchUpdateStores[type_][id]; !ok {
@ -50,7 +50,7 @@ func batchUpdate() {
for i := 0; i < BatchUpdateTypeCount; i++ { for i := 0; i < BatchUpdateTypeCount; i++ {
batchUpdateLocks[i].Lock() batchUpdateLocks[i].Lock()
store := batchUpdateStores[i] store := batchUpdateStores[i]
batchUpdateStores[i] = make(map[int]int) batchUpdateStores[i] = make(map[int]int64)
batchUpdateLocks[i].Unlock() batchUpdateLocks[i].Unlock()
// TODO: maybe we can combine updates with same key? // TODO: maybe we can combine updates with same key?
for key, value := range store { for key, value := range store {
@ -68,7 +68,7 @@ func batchUpdate() {
case BatchUpdateTypeUsedQuota: case BatchUpdateTypeUsedQuota:
updateUserUsedQuota(key, value) updateUserUsedQuota(key, value)
case BatchUpdateTypeRequestCount: case BatchUpdateTypeRequestCount:
updateUserRequestCount(key, value) updateUserRequestCount(key, int(value))
case BatchUpdateTypeChannelUsedQuota: case BatchUpdateTypeChannelUsedQuota:
updateChannelUsedQuota(key, value) updateChannelUsedQuota(key, value)
} }

10
relay/controller/audio.go
View File

@ -50,14 +50,14 @@ func RelayAudioHelper(c *gin.Context, relayMode int) *relaymodel.ErrorWithStatus
modelRatio := common.GetModelRatio(audioModel) modelRatio := common.GetModelRatio(audioModel)
groupRatio := common.GetGroupRatio(group) groupRatio := common.GetGroupRatio(group)
ratio := modelRatio * groupRatio ratio := modelRatio * groupRatio
var quota int var quota int64
var preConsumedQuota int var preConsumedQuota int64
switch relayMode { switch relayMode {
case constant.RelayModeAudioSpeech: case constant.RelayModeAudioSpeech:
preConsumedQuota = int(float64(len(ttsRequest.Input)) * ratio) preConsumedQuota = int64(float64(len(ttsRequest.Input)) * ratio)
quota = preConsumedQuota quota = preConsumedQuota
default: default:
preConsumedQuota = int(float64(config.PreConsumedQuota) * ratio) preConsumedQuota = int64(float64(config.PreConsumedQuota) * ratio)
} }
userQuota, err := model.CacheGetUserQuota(ctx, userId) userQuota, err := model.CacheGetUserQuota(ctx, userId)
if err != nil { if err != nil {
@ -184,7 +184,7 @@ func RelayAudioHelper(c *gin.Context, relayMode int) *relaymodel.ErrorWithStatus
if err != nil { if err != nil {
return openai.ErrorWrapper(err, "get_text_from_body_err", http.StatusInternalServerError) return openai.ErrorWrapper(err, "get_text_from_body_err", http.StatusInternalServerError)
} }
quota = openai.CountTokenText(text, audioModel) quota = int64(openai.CountTokenText(text, audioModel))
resp.Body = io.NopCloser(bytes.NewBuffer(responseBody)) resp.Body = io.NopCloser(bytes.NewBuffer(responseBody))
} }
if resp.StatusCode != http.StatusOK { if resp.StatusCode != http.StatusOK {

14
relay/controller/helper.go
View File

@ -107,15 +107,15 @@ func getPromptTokens(textRequest *relaymodel.GeneralOpenAIRequest, relayMode int
return 0 return 0
} }
func getPreConsumedQuota(textRequest *relaymodel.GeneralOpenAIRequest, promptTokens int, ratio float64) int { func getPreConsumedQuota(textRequest *relaymodel.GeneralOpenAIRequest, promptTokens int, ratio float64) int64 {
preConsumedTokens := config.PreConsumedQuota preConsumedTokens := config.PreConsumedQuota
if textRequest.MaxTokens != 0 { if textRequest.MaxTokens != 0 {
preConsumedTokens = promptTokens + textRequest.MaxTokens preConsumedTokens = int64(promptTokens) + int64(textRequest.MaxTokens)
} }
return int(float64(preConsumedTokens) * ratio) return int64(float64(preConsumedTokens) * ratio)
} }
func preConsumeQuota(ctx context.Context, textRequest *relaymodel.GeneralOpenAIRequest, promptTokens int, ratio float64, meta *util.RelayMeta) (int, *relaymodel.ErrorWithStatusCode) { func preConsumeQuota(ctx context.Context, textRequest *relaymodel.GeneralOpenAIRequest, promptTokens int, ratio float64, meta *util.RelayMeta) (int64, *relaymodel.ErrorWithStatusCode) {
preConsumedQuota := getPreConsumedQuota(textRequest, promptTokens, ratio) preConsumedQuota := getPreConsumedQuota(textRequest, promptTokens, ratio)
userQuota, err := model.CacheGetUserQuota(ctx, meta.UserId) userQuota, err := model.CacheGetUserQuota(ctx, meta.UserId)
@ -144,16 +144,16 @@ func preConsumeQuota(ctx context.Context, textRequest *relaymodel.GeneralOpenAIR
return preConsumedQuota, nil return preConsumedQuota, nil
} }
func postConsumeQuota(ctx context.Context, usage *relaymodel.Usage, meta *util.RelayMeta, textRequest *relaymodel.GeneralOpenAIRequest, ratio float64, preConsumedQuota int, modelRatio float64, groupRatio float64) { func postConsumeQuota(ctx context.Context, usage *relaymodel.Usage, meta *util.RelayMeta, textRequest *relaymodel.GeneralOpenAIRequest, ratio float64, preConsumedQuota int64, modelRatio float64, groupRatio float64) {
if usage == nil { if usage == nil {
logger.Error(ctx, "usage is nil, which is unexpected") logger.Error(ctx, "usage is nil, which is unexpected")
return return
} }
quota := 0 var quota int64
completionRatio := common.GetCompletionRatio(textRequest.Model) completionRatio := common.GetCompletionRatio(textRequest.Model)
promptTokens := usage.PromptTokens promptTokens := usage.PromptTokens
completionTokens := usage.CompletionTokens completionTokens := usage.CompletionTokens
quota = int(math.Ceil((float64(promptTokens) + float64(completionTokens)*completionRatio) * ratio)) quota = int64(math.Ceil((float64(promptTokens) + float64(completionTokens)*completionRatio) * ratio))
if ratio != 0 && quota <= 0 { if ratio != 0 && quota <= 0 {
quota = 1 quota = 1
} }

2
relay/controller/image.go
View File

@ -81,7 +81,7 @@ func RelayImageHelper(c *gin.Context, relayMode int) *relaymodel.ErrorWithStatus
ratio := modelRatio * groupRatio ratio := modelRatio * groupRatio
userQuota, err := model.CacheGetUserQuota(ctx, meta.UserId) userQuota, err := model.CacheGetUserQuota(ctx, meta.UserId)
quota := int(ratio*imageCostRatio*1000) * imageRequest.N quota := int64(ratio*imageCostRatio*1000) * int64(imageRequest.N)
if userQuota-quota < 0 { if userQuota-quota < 0 {
return openai.ErrorWrapper(errors.New("user quota is not enough"), "insufficient_user_quota", http.StatusForbidden) return openai.ErrorWrapper(errors.New("user quota is not enough"), "insufficient_user_quota", http.StatusForbidden)

2
relay/util/billing.go
View File

@ -6,7 +6,7 @@ import (
"github.com/songquanpeng/one-api/model" "github.com/songquanpeng/one-api/model"
) )
func ReturnPreConsumedQuota(ctx context.Context, preConsumedQuota int, tokenId int) { func ReturnPreConsumedQuota(ctx context.Context, preConsumedQuota int64, tokenId int) {
if preConsumedQuota != 0 { if preConsumedQuota != 0 {
go func(ctx context.Context) { go func(ctx context.Context) {
// return pre-consumed quota // return pre-consumed quota

4
relay/util/common.go
View File

@ -155,7 +155,7 @@ func GetFullRequestURL(baseURL string, requestURL string, channelType int) strin
return fullRequestURL return fullRequestURL
} }
func PostConsumeQuota(ctx context.Context, tokenId int, quotaDelta int, totalQuota int, userId int, channelId int, modelRatio float64, groupRatio float64, modelName string, tokenName string) { func PostConsumeQuota(ctx context.Context, tokenId int, quotaDelta int64, totalQuota int64, userId int, channelId int, modelRatio float64, groupRatio float64, modelName string, tokenName string) {
// quotaDelta is remaining quota to be consumed // quotaDelta is remaining quota to be consumed
err := model.PostConsumeTokenQuota(tokenId, quotaDelta) err := model.PostConsumeTokenQuota(tokenId, quotaDelta)
if err != nil { if err != nil {
@ -168,7 +168,7 @@ func PostConsumeQuota(ctx context.Context, tokenId int, quotaDelta int, totalQuo
// totalQuota is total quota consumed // totalQuota is total quota consumed
if totalQuota != 0 { if totalQuota != 0 {
logContent := fmt.Sprintf("模型倍率 %.2f，分组倍率 %.2f", modelRatio, groupRatio) logContent := fmt.Sprintf("模型倍率 %.2f，分组倍率 %.2f", modelRatio, groupRatio)
model.RecordConsumeLog(ctx, userId, channelId, totalQuota, 0, modelName, tokenName, totalQuota, logContent) model.RecordConsumeLog(ctx, userId, channelId, int(totalQuota), 0, modelName, tokenName, totalQuota, logContent)
model.UpdateUserUsedQuotaAndRequestCount(userId, totalQuota) model.UpdateUserUsedQuotaAndRequestCount(userId, totalQuota)
model.UpdateChannelUsedQuota(channelId, totalQuota) model.UpdateChannelUsedQuota(channelId, totalQuota)
} }