feat: able to set pre consumed quota now

README.md

@@ -40,6 +40,9 @@ _✨ All in one 的 OpenAI 接口，整合各种 API 访问方式，开箱即用
   <a href="https://openai.justsong.cn/">在线演示</a>
 </p>
 
+> **Warning**：从 `v0.2` 版本升级到 `v0.3` 版本需要手动迁移数据库，请手动执行[数据库迁移脚本](./bin/migration_v0.2-v0.3.sql)。
+
+
 ## 功能
 1. 支持多种 API 访问渠道，欢迎 PR 或提 issue 添加更多渠道：
    + [x] OpenAI 官方通道

bin/migration_v0.2-v0.3.sql

@@ -0,0 +1,6 @@
+UPDATE users
+SET quota = quota + (
+    SELECT SUM(remain_quota)
+    FROM tokens
+    WHERE tokens.user_id = users.id
+)

common/constants.go

@@ -54,6 +54,7 @@ var QuotaForNewUser = 0
 var ChannelDisableThreshold = 5.0
 var AutomaticDisableChannelEnabled = false
 var QuotaRemindThreshold = 1000
+var PreConsumedQuota = 500
 
 var RootUserEmail = ""
 

controller/relay.go

@@ -128,6 +128,14 @@ func relayHelper(c *gin.Context) error {
 		model_ = strings.TrimSuffix(model_, "-0314")
 		fullRequestURL = fmt.Sprintf("%s/openai/deployments/%s/%s", baseURL, model_, task)
 	}
+	ratio := common.GetModelRatio(textRequest.Model)
+	preConsumedQuota := int(float64(common.PreConsumedQuota) * ratio)
+	if consumeQuota {
+		err := model.PreConsumeTokenQuota(tokenId, preConsumedQuota)
+		if err != nil {
+			return err
+		}
+	}
 	req, err := http.NewRequest(c.Request.Method, fullRequestURL, c.Request.Body)
 	if err != nil {
 		return err
@@ -177,9 +185,9 @@ func relayHelper(c *gin.Context) error {
 			} else {
 				quota = textResponse.Usage.PromptTokens + textResponse.Usage.CompletionTokens*completionRatio
 			}
-			ratio := common.GetModelRatio(textRequest.Model)
 			quota = int(float64(quota) * ratio)
-			err := model.DecreaseTokenQuota(tokenId, quota)
+			quotaDelta := quota - preConsumedQuota
+			err := model.PostConsumeTokenQuota(tokenId, quotaDelta)
 			if err != nil {
 				common.SysError("Error consuming token remain quota: " + err.Error())
 			}

controller/token.go

@@ -160,7 +160,6 @@ func DeleteToken(c *gin.Context) {
 }
 
 func UpdateToken(c *gin.Context) {
-	isAdmin := c.GetInt("role") >= common.RoleAdminUser
 	userId := c.GetInt("id")
 	statusOnly := c.Query("status_only")
 	token := model.Token{}
@@ -191,7 +190,7 @@ func UpdateToken(c *gin.Context) {
 		if cleanToken.Status == common.TokenStatusExhausted && cleanToken.RemainQuota <= 0 && !cleanToken.UnlimitedQuota {
 			c.JSON(http.StatusOK, gin.H{
 				"success": false,
-				"message": "令牌可用次数已用尽，无法启用，请先修改令牌剩余次数，或者设置为无限次数",
+				"message": "令牌可用额度已用尽，无法启用，请先修改令牌剩余额度，或者设置为无限额度",
 			})
 			return
 		}
@@ -202,10 +201,8 @@ func UpdateToken(c *gin.Context) {
 		// If you add more fields, please also update token.Update()
 		cleanToken.Name = token.Name
 		cleanToken.ExpiredTime = token.ExpiredTime
-		if isAdmin {
-			cleanToken.RemainQuota = token.RemainQuota
-			cleanToken.UnlimitedQuota = token.UnlimitedQuota
-		}
+		cleanToken.RemainQuota = token.RemainQuota
+		cleanToken.UnlimitedQuota = token.UnlimitedQuota
 	}
 	err = cleanToken.Update()
 	if err != nil {

middleware/auth.go

@@ -111,7 +111,7 @@ func TokenAuth() func(c *gin.Context) {
 		c.Set("id", token.UserId)
 		c.Set("token_id", token.Id)
 		requestURL := c.Request.URL.String()
-		consumeQuota := !token.UnlimitedQuota
+		consumeQuota := true
 		if strings.HasPrefix(requestURL, "/v1/models") {
 			consumeQuota = false
 		}

model/option.go

@@ -55,6 +55,7 @@ func InitOptionMap() {
 	common.OptionMap["TurnstileSecretKey"] = ""
 	common.OptionMap["QuotaForNewUser"] = strconv.Itoa(common.QuotaForNewUser)
 	common.OptionMap["QuotaRemindThreshold"] = strconv.Itoa(common.QuotaRemindThreshold)
+	common.OptionMap["PreConsumedQuota"] = strconv.Itoa(common.PreConsumedQuota)
 	common.OptionMap["ModelRatio"] = common.ModelRatio2JSONString()
 	common.OptionMap["TopUpLink"] = common.TopUpLink
 	common.OptionMapRWMutex.Unlock()
@@ -159,6 +160,8 @@ func updateOptionMap(key string, value string) (err error) {
 		common.QuotaForNewUser, _ = strconv.Atoi(value)
 	case "QuotaRemindThreshold":
 		common.QuotaRemindThreshold, _ = strconv.Atoi(value)
+	case "PreConsumedQuota":
+		common.PreConsumedQuota, _ = strconv.Atoi(value)
 	case "ModelRatio":
 		err = common.UpdateModelRatioByJSONString(value)
 	case "TopUpLink":

model/token.go

@@ -130,7 +130,23 @@ func DeleteTokenById(id int, userId int) (err error) {
 	return token.Delete()
 }
 
-func DecreaseTokenQuota(tokenId int, quota int) (err error) {
+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
+	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
+	return err
+}
+
+func PreConsumeTokenQuota(tokenId int, quota int) (err error) {
 	if quota < 0 {
 		return errors.New("quota 不能为负数！")
 	}
@@ -138,7 +154,7 @@ func DecreaseTokenQuota(tokenId int, quota int) (err error) {
 	if err != nil {
 		return err
 	}
-	if token.RemainQuota < quota {
+	if !token.UnlimitedQuota && token.RemainQuota < quota {
 		return errors.New("令牌额度不足")
 	}
 	userQuota, err := GetUserQuota(token.UserId)
@@ -163,17 +179,42 @@ func DecreaseTokenQuota(tokenId int, quota int) (err error) {
 			if email != "" {
 				topUpLink := fmt.Sprintf("%s/topup", common.ServerAddress)
 				err = common.SendEmail(prompt, email,
-					fmt.Sprintf("%s，剩余额度为 %d，为了不影响您的使用，请及时充值。<br/>充值链接：<a href='%s'>%s</a>", prompt, userQuota-quota, topUpLink, topUpLink))
+					fmt.Sprintf("%s，当前剩余额度为 %d，为了不影响您的使用，请及时充值。<br/>充值链接：<a href='%s'>%s</a>", prompt, userQuota, topUpLink, topUpLink))
 				if err != nil {
 					common.SysError("发送邮件失败：" + err.Error())
 				}
 			}
 		}()
 	}
-	err = DB.Model(&Token{}).Where("id = ?", tokenId).Update("remain_quota", gorm.Expr("remain_quota - ?", quota)).Error
-	if err != nil {
-		return err
+	if !token.UnlimitedQuota {
+		err = DecreaseTokenQuota(tokenId, quota)
+		if err != nil {
+			return err
+		}
 	}
 	err = DecreaseUserQuota(token.UserId, quota)
 	return err
 }
+
+func PostConsumeTokenQuota(tokenId int, quota int) (err error) {
+	token, err := GetTokenById(tokenId)
+	if quota > 0 {
+		err = DecreaseUserQuota(token.UserId, quota)
+	} else {
+		err = IncreaseUserQuota(token.UserId, -quota)
+	}
+	if err != nil {
+		return err
+	}
+	if !token.UnlimitedQuota {
+		if quota > 0 {
+			err = DecreaseTokenQuota(tokenId, quota)
+		} else {
+			err = IncreaseTokenQuota(tokenId, -quota)
+		}
+		if err != nil {
+			return err
+		}
+	}
+	return nil
+}

web/src/components/Header.js

@@ -34,7 +34,6 @@ const headerButtons = [
     name: '充值',
     to: '/topup',
     icon: 'cart',
-    admin: true,
   },
   {
     name: '用户',

web/src/components/SystemSetting.js

@@ -28,6 +28,7 @@ const SystemSetting = () => {
     RegisterEnabled: '',
     QuotaForNewUser: 0,
     QuotaRemindThreshold: 0,
+    PreConsumedQuota: 0,
     ModelRatio: '',
     TopUpLink: '',
     AutomaticDisableChannelEnabled: '',
@@ -98,6 +99,7 @@ const SystemSetting = () => {
       name === 'TurnstileSecretKey' ||
       name === 'QuotaForNewUser' ||
       name === 'QuotaRemindThreshold' ||
+      name === 'PreConsumedQuota' ||
       name === 'ModelRatio' ||
       name === 'TopUpLink'
     ) {
@@ -119,6 +121,9 @@ const SystemSetting = () => {
     if (originInputs['QuotaRemindThreshold'] !== inputs.QuotaRemindThreshold) {
       await updateOption('QuotaRemindThreshold', inputs.QuotaRemindThreshold);
     }
+    if (originInputs['PreConsumedQuota'] !== inputs.PreConsumedQuota) {
+      await updateOption('PreConsumedQuota', inputs.PreConsumedQuota);
+    }
     if (originInputs['ModelRatio'] !== inputs.ModelRatio) {
       if (!verifyJSON(inputs.ModelRatio)) {
         showError('模型倍率不是合法的 JSON 字符串');
@@ -272,7 +277,7 @@ const SystemSetting = () => {
           <Header as='h3'>
             运营设置
           </Header>
-          <Form.Group widths={3}>
+          <Form.Group widths={4}>
             <Form.Input
               label='新用户初始配额'
               name='QuotaForNewUser'
@@ -302,6 +307,16 @@ const SystemSetting = () => {
               min='0'
               placeholder='低于此额度时将发送邮件提醒用户'
             />
+            <Form.Input
+              label='请求预扣费额度'
+              name='PreConsumedQuota'
+              onChange={handleInputChange}
+              autoComplete='new-password'
+              value={inputs.PreConsumedQuota}
+              type='number'
+              min='0'
+              placeholder='请求结束后多退少补'
+            />
           </Form.Group>
           <Form.Group widths='equal'>
             <Form.TextArea

web/src/pages/Token/EditToken.js

@@ -1,5 +1,5 @@
 import React, { useEffect, useState } from 'react';
-import { Button, Form, Header, Segment } from 'semantic-ui-react';
+import { Button, Form, Header, Message, Segment } from 'semantic-ui-react';
 import { useParams } from 'react-router-dom';
 import { API, showError, showSuccess, timestamp2string } from '../../helpers';
 
@@ -106,6 +106,7 @@ const EditToken = () => {
               required={!isEdit}
             />
           </Form.Field>
+          <Message>注意，令牌的额度仅用于限制令牌本身的最大额度使用量，实际的使用受到账户的剩余额度限制。</Message>
           <Form.Field>
             <Form.Input
               label='额度'