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) { func Relay(c *gin.Context) {
channelType := c.GetInt("channel") channelType := c.GetInt("channel")
tokenId := c.GetInt("token_id") tokenId := c.GetInt("token_id")
isUnlimitedTimes := c.GetBool("unlimited_times") isUnlimitedQuota := c.GetBool("unlimited_quota")
baseURL := common.ChannelBaseURLs[channelType] baseURL := common.ChannelBaseURLs[channelType]
if channelType == common.ChannelTypeCustom { if channelType == common.ChannelTypeCustom {
baseURL = c.GetString("base_url") baseURL = c.GetString("base_url")
@ -56,8 +56,8 @@ func Relay(c *gin.Context) {
if err != nil { if err != nil {
common.SysError("Error closing request body: " + err.Error()) common.SysError("Error closing request body: " + err.Error())
} }
if !isUnlimitedTimes && requestURL == "/v1/chat/completions" { if !isUnlimitedQuota && requestURL == "/v1/chat/completions" {
err := model.DecreaseTokenRemainTimesById(tokenId) err := model.DecreaseTokenRemainQuotaById(tokenId)
if err != nil { if err != nil {
common.SysError("Error decreasing token remain times: " + err.Error()) 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, ExpiredTime: token.ExpiredTime,
} }
if isAdmin { if isAdmin {
cleanToken.RemainTimes = token.RemainTimes cleanToken.RemainQuota = token.RemainQuota
cleanToken.UnlimitedTimes = token.UnlimitedTimes cleanToken.UnlimitedQuota = token.UnlimitedQuota
} else { } else {
userId := c.GetInt("id") userId := c.GetInt("id")
quota, err := model.GetUserQuota(userId) quota, err := model.GetUserQuota(userId)
@ -115,7 +115,7 @@ func AddToken(c *gin.Context) {
return return
} }
if quota > 0 { if quota > 0 {
cleanToken.RemainTimes = quota cleanToken.RemainQuota = quota
} }
} }
err = cleanToken.Insert() err = cleanToken.Insert()
@ -128,7 +128,7 @@ func AddToken(c *gin.Context) {
} }
if !isAdmin { if !isAdmin {
// update user quota // 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{ c.JSON(http.StatusOK, gin.H{
"success": true, "success": true,
@ -184,7 +184,7 @@ func UpdateToken(c *gin.Context) {
}) })
return 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{ c.JSON(http.StatusOK, gin.H{
"success": false, "success": false,
"message": "令牌可用次数已用尽，无法启用，请先修改令牌剩余次数，或者设置为无限次数", "message": "令牌可用次数已用尽，无法启用，请先修改令牌剩余次数，或者设置为无限次数",
@ -199,8 +199,8 @@ func UpdateToken(c *gin.Context) {
cleanToken.Name = token.Name cleanToken.Name = token.Name
cleanToken.ExpiredTime = token.ExpiredTime cleanToken.ExpiredTime = token.ExpiredTime
if isAdmin { if isAdmin {
cleanToken.RemainTimes = token.RemainTimes cleanToken.RemainQuota = token.RemainQuota
cleanToken.UnlimitedTimes = token.UnlimitedTimes cleanToken.UnlimitedQuota = token.UnlimitedQuota
} }
} }
err = cleanToken.Update() 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("id", token.UserId)
c.Set("token_id", token.Id) c.Set("token_id", token.Id)
c.Set("unlimited_times", token.UnlimitedTimes) c.Set("unlimited_quota", token.UnlimitedQuota)
if len(parts) > 1 { if len(parts) > 1 {
if model.IsAdmin(token.UserId) { if model.IsAdmin(token.UserId) {
c.Set("channelId", parts[1]) 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"` CreatedTime int64 `json:"created_time" gorm:"bigint"`
AccessedTime int64 `json:"accessed_time" gorm:"bigint"` AccessedTime int64 `json:"accessed_time" gorm:"bigint"`
ExpiredTime int64 `json:"expired_time" gorm:"bigint;default:-1"` // -1 means never expired ExpiredTime int64 `json:"expired_time" gorm:"bigint;default:-1"` // -1 means never expired
RemainTimes int `json:"remain_times" gorm:"default:0"` RemainQuota int `json:"remain_quota" gorm:"default:0"`
UnlimitedTimes bool `json:"unlimited_times" gorm:"default:false"` UnlimitedQuota bool `json:"unlimited_quota" gorm:"default:false"`
} }
func GetAllUserTokens(userId int, startIdx int, num int) ([]*Token, error) { 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 已过期") return nil, errors.New("该 token 已过期")
} }
if !token.UnlimitedTimes && token.RemainTimes <= 0 { if !token.UnlimitedQuota && token.RemainQuota <= 0 {
token.Status = common.TokenStatusExhausted token.Status = common.TokenStatusExhausted
err := token.SelectUpdate() err := token.SelectUpdate()
if err != nil { if err != nil {
common.SysError("更新 token 状态失败：" + err.Error()) common.SysError("更新 token 状态失败：" + err.Error())
} }
return nil, errors.New("该 token 可用次数已用尽") return nil, errors.New("该 token 额度已用尽")
} }
go func() { go func() {
token.AccessedTime = common.GetTimestamp() 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 // Update Make sure your token's fields is completed, because this will update non-zero values
func (token *Token) Update() error { func (token *Token) Update() error {
var err 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 return err
} }
@ -119,12 +119,12 @@ func DeleteTokenById(id int, userId int) (err error) {
return token.Delete() return token.Delete()
} }
func DecreaseTokenRemainTimesById(id int) (err error) { func DecreaseTokenRemainQuotaById(id int) (err error) {
err = DB.Model(&Token{}).Where("id = ?", id).Update("remain_times", gorm.Expr("remain_times - ?", 1)).Error err = DB.Model(&Token{}).Where("id = ?", id).Update("remain_quota", gorm.Expr("remain_quota - ?", 1)).Error
return err return err
} }
func TopUpToken(id int, times int) (err error) { 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 return err
} }

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

@ -164,7 +164,7 @@ const TokensTable = () => {
showSuccess('充值成功！'); showSuccess('充值成功！');
let newTokens = [...tokens]; let newTokens = [...tokens];
let realIdx = (activePage - 1) * ITEMS_PER_PAGE + targetTokenIdx; let realIdx = (activePage - 1) * ITEMS_PER_PAGE + targetTokenIdx;
newTokens[realIdx].remain_times += data; newTokens[realIdx].remain_quota += data;
setTokens(newTokens); setTokens(newTokens);
setRedemptionCode(''); setRedemptionCode('');
setShowTopUpModal(false); setShowTopUpModal(false);
@ -217,10 +217,10 @@ const TokensTable = () => {
<Table.HeaderCell <Table.HeaderCell
style={{ cursor: 'pointer' }} style={{ cursor: 'pointer' }}
onClick={() => { onClick={() => {
sortToken('remain_times'); sortToken('remain_quota');
}} }}
> >
剩余次数 额度
</Table.HeaderCell> </Table.HeaderCell>
<Table.HeaderCell <Table.HeaderCell
style={{ cursor: 'pointer' }} style={{ cursor: 'pointer' }}
@ -255,7 +255,7 @@ const TokensTable = () => {
<Table.Cell>{token.id}</Table.Cell> <Table.Cell>{token.id}</Table.Cell>
<Table.Cell>{token.name ? token.name : '无'}</Table.Cell> <Table.Cell>{token.name ? token.name : '无'}</Table.Cell>
<Table.Cell>{renderStatus(token.status)}</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>{renderTimestamp(token.created_time)}</Table.Cell>
<Table.Cell>{token.expired_time === -1 ? '永不过期' : renderTimestamp(token.expired_time)}</Table.Cell> <Table.Cell>{token.expired_time === -1 ? '永不过期' : renderTimestamp(token.expired_time)}</Table.Cell>
<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 [loading, setLoading] = useState(isEdit);
const originInputs = { const originInputs = {
name: '', name: '',
remain_times: 0, remain_quota: 0,
expired_time: -1, expired_time: -1,
unlimited_times: false unlimited_quota: false
}; };
const isAdminUser = isAdmin(); const isAdminUser = isAdmin();
const [inputs, setInputs] = useState(originInputs); 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 }) => { const handleInputChange = (e, { name, value }) => {
setInputs((inputs) => ({ ...inputs, [name]: value })); setInputs((inputs) => ({ ...inputs, [name]: value }));
@ -37,8 +37,8 @@ const EditToken = () => {
} }
}; };
const setUnlimitedTimes = () => { const setUnlimitedQuota = () => {
setInputs({ ...inputs, unlimited_times: !unlimited_times }); setInputs({ ...inputs, unlimited_quota: !unlimited_quota });
}; };
const loadToken = async () => { const loadToken = async () => {
@ -63,7 +63,7 @@ const EditToken = () => {
const submit = async () => { const submit = async () => {
if (!isEdit && inputs.name === '') return; if (!isEdit && inputs.name === '') return;
let localInputs = inputs; let localInputs = inputs;
localInputs.remain_times = parseInt(localInputs.remain_times); localInputs.remain_quota = parseInt(localInputs.remain_quota);
if (localInputs.expired_time !== -1) { if (localInputs.expired_time !== -1) {
let time = Date.parse(localInputs.expired_time); let time = Date.parse(localInputs.expired_time);
if (isNaN(time)) { if (isNaN(time)) {
@ -111,19 +111,19 @@ const EditToken = () => {
isAdminUser && <> isAdminUser && <>
<Form.Field> <Form.Field>
<Form.Input <Form.Input
label='剩余次数' label='额度'
name='remain_times' name='remain_quota'
placeholder={'请输入剩余次数'} placeholder={'请输入额度'}
onChange={handleInputChange} onChange={handleInputChange}
value={remain_times} value={remain_quota}
autoComplete='off' autoComplete='off'
type='number' type='number'
disabled={unlimited_times} disabled={unlimited_quota}
/> />
</Form.Field> </Form.Field>
<Button type={'button'} style={{marginBottom: '14px'}} onClick={() => { <Button type={'button'} style={{marginBottom: '14px'}} onClick={() => {
setUnlimitedTimes(); setUnlimitedQuota();
}}>{unlimited_times ? '取消无限次' : '设置为无限次'}</Button> }}>{unlimited_quota ? '取消无限额度' : '设置为无限额度'}</Button>
</> </>
} }
<Form.Field> <Form.Field>