From 7edc2b53762dcf5d63878864e3d9842926481812 Mon Sep 17 00:00:00 2001 From: JustSong Date: Fri, 23 Jun 2023 20:14:53 +0800 Subject: [PATCH] feat: able to display token billing stat via billing api (close #186) --- common/constants.go | 3 ++- controller/billing.go | 26 ++++++++++++++++++++++---- controller/option.go | 2 +- model/option.go | 3 +++ model/token.go | 15 +++++++++++++-- web/src/components/OperationSetting.js | 9 ++++++++- 6 files changed, 49 insertions(+), 9 deletions(-) diff --git a/common/constants.go b/common/constants.go index 4439e2a5..04502f28 100644 --- a/common/constants.go +++ b/common/constants.go @@ -18,7 +18,8 @@ var Logo = "" var TopUpLink = "" var ChatLink = "" var QuotaPerUnit = 500 * 1000.0 // $0.002 / 1K tokens -var DisplayInCurrencyEnabled = false +var DisplayInCurrencyEnabled = true +var DisplayTokenStatEnabled = true var UsingSQLite = false diff --git a/controller/billing.go b/controller/billing.go index 8197bc88..ec8e3ce3 100644 --- a/controller/billing.go +++ b/controller/billing.go @@ -7,8 +7,17 @@ import ( ) func GetSubscription(c *gin.Context) { - userId := c.GetInt("id") - quota, err := model.GetUserQuota(userId) + var quota int + var err error + var token *model.Token + if common.DisplayTokenStatEnabled { + tokenId := c.GetInt("token_id") + token, err = model.GetTokenById(tokenId) + quota = token.RemainQuota + } else { + userId := c.GetInt("id") + quota, err = model.GetUserQuota(userId) + } if err != nil { openAIError := OpenAIError{ Message: err.Error(), @@ -35,8 +44,17 @@ func GetSubscription(c *gin.Context) { } func GetUsage(c *gin.Context) { - userId := c.GetInt("id") - quota, err := model.GetUserUsedQuota(userId) + var quota int + var err error + var token *model.Token + if common.DisplayTokenStatEnabled { + tokenId := c.GetInt("token_id") + token, err = model.GetTokenById(tokenId) + quota = token.UsedQuota + } else { + userId := c.GetInt("id") + quota, err = model.GetUserUsedQuota(userId) + } if err != nil { openAIError := OpenAIError{ Message: err.Error(), diff --git a/controller/option.go b/controller/option.go index b5b675c6..abf0d5be 100644 --- a/controller/option.go +++ b/controller/option.go @@ -13,7 +13,7 @@ func GetOptions(c *gin.Context) { var options []*model.Option common.OptionMapRWMutex.Lock() for k, v := range common.OptionMap { - if strings.Contains(k, "Token") || strings.Contains(k, "Secret") { + if strings.HasSuffix(k, "Token") || strings.HasSuffix(k, "Secret") { continue } options = append(options, &model.Option{ diff --git a/model/option.go b/model/option.go index cbf14d0b..f41cf954 100644 --- a/model/option.go +++ b/model/option.go @@ -36,6 +36,7 @@ func InitOptionMap() { common.OptionMap["AutomaticDisableChannelEnabled"] = strconv.FormatBool(common.AutomaticDisableChannelEnabled) common.OptionMap["LogConsumeEnabled"] = strconv.FormatBool(common.LogConsumeEnabled) common.OptionMap["DisplayInCurrencyEnabled"] = strconv.FormatBool(common.DisplayInCurrencyEnabled) + common.OptionMap["DisplayTokenStatEnabled"] = strconv.FormatBool(common.DisplayTokenStatEnabled) common.OptionMap["ChannelDisableThreshold"] = strconv.FormatFloat(common.ChannelDisableThreshold, 'f', -1, 64) common.OptionMap["SMTPServer"] = "" common.OptionMap["SMTPFrom"] = "" @@ -144,6 +145,8 @@ func updateOptionMap(key string, value string) (err error) { common.LogConsumeEnabled = boolValue case "DisplayInCurrencyEnabled": common.DisplayInCurrencyEnabled = boolValue + case "DisplayTokenStatEnabled": + common.DisplayTokenStatEnabled = boolValue } } switch key { diff --git a/model/token.go b/model/token.go index 83760f45..7cd226c6 100644 --- a/model/token.go +++ b/model/token.go @@ -18,6 +18,7 @@ type Token struct { ExpiredTime int64 `json:"expired_time" gorm:"bigint;default:-1"` // -1 means never expired RemainQuota int `json:"remain_quota" gorm:"default:0"` UnlimitedQuota bool `json:"unlimited_quota" gorm:"default:false"` + UsedQuota int `json:"used_quota" gorm:"default:0"` // used quota } func GetAllUserTokens(userId int, startIdx int, num int) ([]*Token, error) { @@ -130,7 +131,12 @@ func IncreaseTokenQuota(id int, quota int) (err error) { if quota < 0 { return errors.New("quota 不能为负数！") } - err = DB.Model(&Token{}).Where("id = ?", id).Update("remain_quota", gorm.Expr("remain_quota + ?", quota)).Error + err = DB.Model(&Token{}).Where("id = ?", id).Updates( + map[string]interface{}{ + "remain_quota": gorm.Expr("remain_quota + ?", quota), + "used_quota": gorm.Expr("used_quota - ?", quota), + }, + ).Error return err } @@ -138,7 +144,12 @@ func DecreaseTokenQuota(id int, quota int) (err error) { if quota < 0 { return errors.New("quota 不能为负数！") } - err = DB.Model(&Token{}).Where("id = ?", id).Update("remain_quota", gorm.Expr("remain_quota - ?", quota)).Error + err = DB.Model(&Token{}).Where("id = ?", id).Updates( + map[string]interface{}{ + "remain_quota": gorm.Expr("remain_quota - ?", quota), + "used_quota": gorm.Expr("used_quota + ?", quota), + }, + ).Error return err } diff --git a/web/src/components/OperationSetting.js b/web/src/components/OperationSetting.js index d8c41230..6de15f16 100644 --- a/web/src/components/OperationSetting.js +++ b/web/src/components/OperationSetting.js @@ -17,7 +17,8 @@ const OperationSetting = () => { AutomaticDisableChannelEnabled: '', ChannelDisableThreshold: 0, LogConsumeEnabled: '', - DisplayInCurrencyEnabled: '' + DisplayInCurrencyEnabled: '', + DisplayTokenStatEnabled: '' }); const [originInputs, setOriginInputs] = useState({}); let [loading, setLoading] = useState(false); @@ -177,6 +178,12 @@ const OperationSetting = () => { name='DisplayInCurrencyEnabled' onChange={handleInputChange} /> + { submitConfig('general').then();