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

fix: return quota to user when delete token (close #37)

Browse Source
This commit is contained in:
JustSong 2023-05-04 10:20:39 +08:00
parent 4fed003f1a
commit 331177d97e
4 changed files with 24 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
controller/relay.go
View File

@ -137,7 +137,7 @@ func relayHelper(c *gin.Context) error {
ratio = common.RatioGPT3dot5 ratio = common.RatioGPT3dot5
} }
quota = int(float64(quota) * ratio) quota = int(float64(quota) * ratio)
err := model.ConsumeTokenQuota(tokenId, quota) err := model.DecreaseTokenQuota(tokenId, quota)
if err != nil { if err != nil {
common.SysError("Error consuming token remain quota: " + err.Error()) common.SysError("Error consuming token remain quota: " + err.Error())
} }

2
model/redemption.go
View File

@ -55,7 +55,7 @@ func Redeem(key string, tokenId int) (quota int, err error) {
if redemption.Status != common.RedemptionCodeStatusEnabled { if redemption.Status != common.RedemptionCodeStatusEnabled {
return 0, errors.New("该兑换码已被使用") return 0, errors.New("该兑换码已被使用")
} }
err = TopUpTokenQuota(tokenId, redemption.Quota) err = IncreaseTokenQuota(tokenId, redemption.Quota)
if err != nil { if err != nil {
return 0, err return 0, err
} }

23
model/token.go
View File

@ -116,15 +116,26 @@ func DeleteTokenById(id int, userId int) (err error) {
if err != nil { if err != nil {
return err return err
} }
quota := token.RemainQuota
if quota != 0 {
if quota > 0 {
err = IncreaseUserQuota(userId, quota)
} else {
err = DecreaseUserQuota(userId, quota)
}
}
if err != nil {
return err
}
return token.Delete() return token.Delete()
} }
func ConsumeTokenQuota(id int, quota int) (err error) { func IncreaseTokenQuota(id int, quota int) (err error) {
err = DB.Model(&Token{}).Where("id = ?", id).Update("remain_quota", gorm.Expr("remain_quota - ?", quota)).Error
return err
}
func TopUpTokenQuota(id int, quota int) (err error) {
err = DB.Model(&Token{}).Where("id = ?", id).Update("remain_quota", gorm.Expr("remain_quota + ?", quota)).Error err = DB.Model(&Token{}).Where("id = ?", id).Update("remain_quota", gorm.Expr("remain_quota + ?", quota)).Error
return err return err
} }
func DecreaseTokenQuota(id int, quota int) (err error) {
err = DB.Model(&Token{}).Where("id = ?", id).Update("remain_quota", gorm.Expr("remain_quota - ?", quota)).Error
return err
}

5
model/user.go
View File

@ -225,6 +225,11 @@ func GetUserQuota(id int) (quota int, err error) {
return quota, err return quota, err
} }
func IncreaseUserQuota(id int, quota int) (err error) {
err = DB.Model(&User{}).Where("id = ?", id).Update("quota", gorm.Expr("quota + ?", quota)).Error
return err
}
func DecreaseUserQuota(id int, quota int) (err error) { func DecreaseUserQuota(id int, quota int) (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