feat: able to display token billing stat via billing api (close #186)

songquanpeng · Jun 23, 2023 · 7edc2b5 · 7edc2b5
1 parent d4869df
commit 7edc2b5
Show file tree

Hide file tree

Showing 6 changed files with 49 additions and 9 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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,15 +131,25 @@ 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
 }
 
 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
@@ -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}
             />
+            <Form.Checkbox
+              checked={inputs.DisplayTokenStatEnabled === 'true'}
+              label='Billing 相关 API 显示令牌额度而非用户额度'
+              name='DisplayTokenStatEnabled'
+              onChange={handleInputChange}
+            />
           </Form.Group>
           <Form.Button onClick={() => {
             submitConfig('general').then();