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

refactor: use quota instead of times

Browse Source
This commit is contained in:
JustSong 2023-04-28 14:57:20 +08:00
parent 7a5057f02d
commit 601fa5cea8
6 changed files with 36 additions and 36 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
controller/relay.go
View File

@ -14,7 +14,7 @@ import (
func Relay(c *gin.Context) {
channelType := c.GetInt("channel")
tokenId := c.GetInt("token_id")
isUnlimitedTimes := c.GetBool("unlimited_times")
isUnlimitedQuota := c.GetBool("unlimited_quota")
baseURL := common.ChannelBaseURLs[channelType]
if channelType == common.ChannelTypeCustom {
baseURL = c.GetString("base_url")
@ -56,8 +56,8 @@ func Relay(c *gin.Context) {
if err != nil {
common.SysError("Error closing request body: " + err.Error())
}
if !isUnlimitedTimes && requestURL == "/v1/chat/completions" {
err := model.DecreaseTokenRemainTimesById(tokenId)
if !isUnlimitedQuota && requestURL == "/v1/chat/completions" {
err := model.DecreaseTokenRemainQuotaById(tokenId)
if err != nil {
common.SysError("Error decreasing token remain times: " + err.Error())
}

14
controller/token.go
View File

@ -102,8 +102,8 @@ func AddToken(c *gin.Context) {
ExpiredTime: token.ExpiredTime,
}
if isAdmin {
cleanToken.RemainTimes = token.RemainTimes
cleanToken.UnlimitedTimes = token.UnlimitedTimes
cleanToken.RemainQuota = token.RemainQuota
cleanToken.UnlimitedQuota = token.UnlimitedQuota
} else {
userId := c.GetInt("id")
quota, err := model.GetUserQuota(userId)
@ -115,7 +115,7 @@ func AddToken(c *gin.Context) {
return
}
if quota > 0 {
cleanToken.RemainTimes = quota
cleanToken.RemainQuota = quota
}
}
err = cleanToken.Insert()
@ -128,7 +128,7 @@ func AddToken(c *gin.Context) {
}
if !isAdmin {
// update user quota
err = model.DecreaseUserQuota(c.GetInt("id"), cleanToken.RemainTimes)
err = model.DecreaseUserQuota(c.GetInt("id"), cleanToken.RemainQuota)
}
c.JSON(http.StatusOK, gin.H{
"success": true,
@ -184,7 +184,7 @@ func UpdateToken(c *gin.Context) {
})
return
}
if cleanToken.Status == common.TokenStatusExhausted && cleanToken.RemainTimes <= 0 && !cleanToken.UnlimitedTimes {
if cleanToken.Status == common.TokenStatusExhausted && cleanToken.RemainQuota <= 0 && !cleanToken.UnlimitedQuota {
c.JSON(http.StatusOK, gin.H{
"success": false,
"message": "令牌可用次数已用尽，无法启用，请先修改令牌剩余次数，或者设置为无限次数",
@ -199,8 +199,8 @@ func UpdateToken(c *gin.Context) {
cleanToken.Name = token.Name
cleanToken.ExpiredTime = token.ExpiredTime
if isAdmin {
cleanToken.RemainTimes = token.RemainTimes
cleanToken.UnlimitedTimes = token.UnlimitedTimes
cleanToken.RemainQuota = token.RemainQuota
cleanToken.UnlimitedQuota = token.UnlimitedQuota
}
}
err = cleanToken.Update()

2
middleware/auth.go
View File

@ -110,7 +110,7 @@ func TokenAuth() func(c *gin.Context) {
}
c.Set("id", token.UserId)
c.Set("token_id", token.Id)
c.Set("unlimited_times", token.UnlimitedTimes)
c.Set("unlimited_quota", token.UnlimitedQuota)
if len(parts) > 1 {
if model.IsAdmin(token.UserId) {
c.Set("channelId", parts[1])

16
model/token.go
View File

@ -17,8 +17,8 @@ type Token struct {
CreatedTime int64 `json:"created_time" gorm:"bigint"`
AccessedTime int64 `json:"accessed_time" gorm:"bigint"`
ExpiredTime int64 `json:"expired_time" gorm:"bigint;default:-1"` // -1 means never expired
RemainTimes int `json:"remain_times" gorm:"default:0"`
UnlimitedTimes bool `json:"unlimited_times" gorm:"default:false"`
RemainQuota int `json:"remain_quota" gorm:"default:0"`
UnlimitedQuota bool `json:"unlimited_quota" gorm:"default:false"`
}
func GetAllUserTokens(userId int, startIdx int, num int) ([]*Token, error) {
@ -52,13 +52,13 @@ func ValidateUserToken(key string) (token *Token, err error) {
}
return nil, errors.New("该 token 已过期")
}
if !token.UnlimitedTimes && token.RemainTimes <= 0 {
if !token.UnlimitedQuota && token.RemainQuota <= 0 {
token.Status = common.TokenStatusExhausted
err := token.SelectUpdate()
if err != nil {
common.SysError("更新 token 状态失败：" + err.Error())
}
return nil, errors.New("该 token 可用次数已用尽")
return nil, errors.New("该 token 额度已用尽")
}
go func() {
token.AccessedTime = common.GetTimestamp()
@ -91,7 +91,7 @@ func (token *Token) Insert() error {
// Update Make sure your token's fields is completed, because this will update non-zero values
func (token *Token) Update() error {
var err error
err = DB.Model(token).Select("name", "status", "expired_time", "remain_times", "unlimited_times").Updates(token).Error
err = DB.Model(token).Select("name", "status", "expired_time", "remain_quota", "unlimited_quota").Updates(token).Error
return err
}
@ -119,12 +119,12 @@ func DeleteTokenById(id int, userId int) (err error) {
return token.Delete()
}
func DecreaseTokenRemainTimesById(id int) (err error) {
err = DB.Model(&Token{}).Where("id = ?", id).Update("remain_times", gorm.Expr("remain_times - ?", 1)).Error
func DecreaseTokenRemainQuotaById(id int) (err error) {
err = DB.Model(&Token{}).Where("id = ?", id).Update("remain_quota", gorm.Expr("remain_quota - ?", 1)).Error
return err
}
func TopUpToken(id int, times int) (err error) {
err = DB.Model(&Token{}).Where("id = ?", id).Update("remain_times", gorm.Expr("remain_times + ?", times)).Error
err = DB.Model(&Token{}).Where("id = ?", id).Update("remain_quota", gorm.Expr("remain_quota + ?", times)).Error
return err
}

8
web/src/components/TokensTable.js
View File

@ -164,7 +164,7 @@ const TokensTable = () => {
showSuccess('充值成功！');
let newTokens = [...tokens];
let realIdx = (activePage - 1) * ITEMS_PER_PAGE + targetTokenIdx;
newTokens[realIdx].remain_times += data;
newTokens[realIdx].remain_quota += data;
setTokens(newTokens);
setRedemptionCode('');
setShowTopUpModal(false);
@ -217,10 +217,10 @@ const TokensTable = () => {
<Table.HeaderCell
style={{ cursor: 'pointer' }}
onClick={() => {
sortToken('remain_times');
sortToken('remain_quota');
}}
>
剩余次数
额度
</Table.HeaderCell>
<Table.HeaderCell
style={{ cursor: 'pointer' }}
@ -255,7 +255,7 @@ const TokensTable = () => {
<Table.Cell>{token.id}</Table.Cell>
<Table.Cell>{token.name ? token.name : '无'}</Table.Cell>
<Table.Cell>{renderStatus(token.status)}</Table.Cell>
<Table.Cell>{token.unlimited_times ? '无限制' : token.remain_times}</Table.Cell>
<Table.Cell>{token.unlimited_quota ? '无限制' : token.remain_quota}</Table.Cell>
<Table.Cell>{renderTimestamp(token.created_time)}</Table.Cell>
<Table.Cell>{token.expired_time === -1 ? '永不过期' : renderTimestamp(token.expired_time)}</Table.Cell>
<Table.Cell>

26
web/src/pages/Token/EditToken.js
View File

@ -10,13 +10,13 @@ const EditToken = () => {
const [loading, setLoading] = useState(isEdit);
const originInputs = {
name: '',
remain_times: 0,
remain_quota: 0,
expired_time: -1,
unlimited_times: false
unlimited_quota: false
};
const isAdminUser = isAdmin();
const [inputs, setInputs] = useState(originInputs);
const { name, remain_times, expired_time, unlimited_times } = inputs;
const { name, remain_quota, expired_time, unlimited_quota } = inputs;
const handleInputChange = (e, { name, value }) => {
setInputs((inputs) => ({ ...inputs, [name]: value }));
@ -37,8 +37,8 @@ const EditToken = () => {
}
};
const setUnlimitedTimes = () => {
setInputs({ ...inputs, unlimited_times: !unlimited_times });
const setUnlimitedQuota = () => {
setInputs({ ...inputs, unlimited_quota: !unlimited_quota });
};
const loadToken = async () => {
@ -63,7 +63,7 @@ const EditToken = () => {
const submit = async () => {
if (!isEdit && inputs.name === '') return;
let localInputs = inputs;
localInputs.remain_times = parseInt(localInputs.remain_times);
localInputs.remain_quota = parseInt(localInputs.remain_quota);
if (localInputs.expired_time !== -1) {
let time = Date.parse(localInputs.expired_time);
if (isNaN(time)) {
@ -111,19 +111,19 @@ const EditToken = () => {
isAdminUser && <>
<Form.Field>
<Form.Input
label='剩余次数'
name='remain_times'
placeholder={'请输入剩余次数'}
label='额度'
name='remain_quota'
placeholder={'请输入额度'}
onChange={handleInputChange}
value={remain_times}
value={remain_quota}
autoComplete='off'
type='number'
disabled={unlimited_times}
disabled={unlimited_quota}
/>
</Form.Field>
<Button type={'button'} style={{marginBottom: '14px'}} onClick={() => {
setUnlimitedTimes();
}}>{unlimited_times ? '取消无限次' : '设置为无限次'}</Button>
setUnlimitedQuota();
}}>{unlimited_quota ? '取消无限额度' : '设置为无限额度'}</Button>
</>
}
<Form.Field>