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

fix: only reduce remain times when request /v1/chat/completions (close #15)

Browse Source 
BREAKING CHANGE: now remain_times is -1 doesn't mean unlimited times anymore!
This commit is contained in:
JustSong 2023-04-26 10:45:34 +08:00
parent eb8f43acb5
commit 109736cc05
6 changed files with 83 additions and 40 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

20
controller/relay.go
View File

@ -7,16 +7,20 @@ import (
"io"
"net/http"
"one-api/common"
"one-api/model"
"strings"
)
func Relay(c *gin.Context) {
channelType := c.GetInt("channel")
tokenId := c.GetInt("token_id")
isUnlimitedTimes := c.GetBool("unlimited_times")
baseURL := common.ChannelBaseURLs[channelType]
if channelType == common.ChannelTypeCustom {
baseURL = c.GetString("base_url")
}
req, err := http.NewRequest(c.Request.Method, fmt.Sprintf("%s%s", baseURL, c.Request.URL.String()), c.Request.Body)
requestURL := c.Request.URL.String()
req, err := http.NewRequest(c.Request.Method, fmt.Sprintf("%s%s", baseURL, requestURL), c.Request.Body)
if err != nil {
c.JSON(http.StatusOK, gin.H{
"error": gin.H{
@ -46,7 +50,19 @@ func Relay(c *gin.Context) {
})
return
}
defer resp.Body.Close()
defer func() {
err := req.Body.Close()
if err != nil {
common.SysError("Error closing request body: " + err.Error())
}
if !isUnlimitedTimes && requestURL == "/v1/chat/completions" {
err := model.DecreaseTokenRemainTimesById(tokenId)
if err != nil {
common.SysError("Error decreasing token remain times: " + err.Error())
}
}
}()
isStream := resp.Header.Get("Content-Type") == "text/event-stream"
if isStream {
scanner := bufio.NewScanner(resp.Body)

34
controller/token.go
View File

@ -93,13 +93,14 @@ func AddToken(c *gin.Context) {
return
}
cleanToken := model.Token{
UserId: c.GetInt("id"),
Name: token.Name,
Key: common.GetUUID(),
CreatedTime: common.GetTimestamp(),
AccessedTime: common.GetTimestamp(),
ExpiredTime: token.ExpiredTime,
RemainTimes: token.RemainTimes,
UserId: c.GetInt("id"),
Name: token.Name,
Key: common.GetUUID(),
CreatedTime: common.GetTimestamp(),
AccessedTime: common.GetTimestamp(),
ExpiredTime: token.ExpiredTime,
RemainTimes: token.RemainTimes,
UnlimitedTimes: token.UnlimitedTimes,
}
err = cleanToken.Insert()
if err != nil {
@ -136,6 +137,7 @@ func DeleteToken(c *gin.Context) {
func UpdateToken(c *gin.Context) {
userId := c.GetInt("id")
statusOnly := c.Query("status_only")
token := model.Token{}
err := c.ShouldBindJSON(&token)
if err != nil {
@ -161,19 +163,23 @@ func UpdateToken(c *gin.Context) {
})
return
}
if cleanToken.Status == common.TokenStatusExhausted && cleanToken.RemainTimes == 0 {
if cleanToken.Status == common.TokenStatusExhausted && cleanToken.RemainTimes <= 0 && !cleanToken.UnlimitedTimes {
c.JSON(http.StatusOK, gin.H{
"success": false,
"message": "令牌可用次数已用尽，无法启用，请先修改令牌剩余次数",
"message": "令牌可用次数已用尽，无法启用，请先修改令牌剩余次数，或者设置为无限次数",
})
return
}
}
cleanToken.Name = token.Name
cleanToken.Status = token.Status
cleanToken.ExpiredTime = token.ExpiredTime
cleanToken.RemainTimes = token.RemainTimes
if statusOnly != "" {
cleanToken.Status = token.Status
} else {
// If you add more fields, please also update token.Update()
cleanToken.Name = token.Name
cleanToken.ExpiredTime = token.ExpiredTime
cleanToken.RemainTimes = token.RemainTimes
cleanToken.UnlimitedTimes = token.UnlimitedTimes
}
err = cleanToken.Update()
if err != nil {
c.JSON(http.StatusOK, gin.H{

2
middleware/auth.go
View File

@ -80,6 +80,8 @@ func TokenAuth() func(c *gin.Context) {
return
}
c.Set("id", token.UserId)
c.Set("token_id", token.Id)
c.Set("unlimited_times", token.UnlimitedTimes)
if len(parts) > 1 {
c.Set("channelId", parts[1])
}

44
model/token.go
View File

@ -3,20 +3,22 @@ package model
import (
"errors"
_ "gorm.io/driver/sqlite"
"gorm.io/gorm"
"one-api/common"
"strings"
)
type Token struct {
Id int `json:"id"`
UserId int `json:"user_id"`
Key string `json:"key" gorm:"uniqueIndex"`
Status int `json:"status" gorm:"default:1"`
Name string `json:"name" gorm:"index" `
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:-1"` // -1 means infinite times
Id int `json:"id"`
UserId int `json:"user_id"`
Key string `json:"key" gorm:"uniqueIndex"`
Status int `json:"status" gorm:"default:1"`
Name string `json:"name" gorm:"index" `
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"`
}
func GetAllUserTokens(userId int, startIdx int, num int) ([]*Token, error) {
@ -50,14 +52,16 @@ func ValidateUserToken(key string) (token *Token, err error) {
}
return nil, errors.New("该 token 已过期")
}
if !token.UnlimitedTimes && token.RemainTimes <= 0 {
token.Status = common.TokenStatusExhausted
err := token.SelectUpdate()
if err != nil {
common.SysError("更新 token 状态失败：" + err.Error())
}
return nil, errors.New("该 token 可用次数已用尽")
}
go func() {
token.AccessedTime = common.GetTimestamp()
if token.RemainTimes > 0 {
token.RemainTimes--
if token.RemainTimes == 0 {
token.Status = common.TokenStatusExhausted
}
}
err := token.SelectUpdate()
if err != nil {
common.SysError("更新 token 失败：" + err.Error())
@ -84,15 +88,16 @@ func (token *Token) Insert() error {
return err
}
// 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).Updates(token).Error
err = DB.Model(token).Select("name", "status", "expired_time", "remain_times", "unlimited_times").Updates(token).Error
return err
}
func (token *Token) SelectUpdate() error {
// This can update zero values
return DB.Model(token).Select("accessed_time", "remain_times", "status").Updates(token).Error
return DB.Model(token).Select("accessed_time", "status").Updates(token).Error
}
func (token *Token) Delete() error {
@ -113,3 +118,8 @@ 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
return err
}

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

@ -79,11 +79,11 @@ const TokensTable = () => {
break;
case 'enable':
data.status = 1;
res = await API.put('/api/token/', data);
res = await API.put('/api/token/?status_only=true', data);
break;
case 'disable':
data.status = 2;
res = await API.put('/api/token/', data);
res = await API.put('/api/token/?status_only=true', data);
break;
}
const { success, message } = res.data;
@ -230,7 +230,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.remain_times === -1 ? "无限制" : token.remain_times}</Table.Cell>
<Table.Cell>{token.unlimited_times ? "无限制" : token.remain_times}</Table.Cell>
<Table.Cell>{renderTimestamp(token.created_time)}</Table.Cell>
<Table.Cell>{renderTimestamp(token.accessed_time)}</Table.Cell>
<Table.Cell>{token.expired_time === -1 ? "永不过期" : renderTimestamp(token.expired_time)}</Table.Cell>

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

@ -10,11 +10,12 @@ const EditToken = () => {
const [loading, setLoading] = useState(isEdit);
const originInputs = {
name: '',
remain_times: -1,
expired_time: -1
remain_times: 0,
expired_time: -1,
unlimited_times: false,
};
const [inputs, setInputs] = useState(originInputs);
const { name, remain_times, expired_time } = inputs;
const { name, remain_times, expired_time, unlimited_times } = inputs;
const handleInputChange = (e, { name, value }) => {
setInputs((inputs) => ({ ...inputs, [name]: value }));
@ -35,6 +36,10 @@ const EditToken = () => {
}
};
const setUnlimitedTimes = () => {
setInputs({ ...inputs, unlimited_times: !unlimited_times });
}
const loadToken = async () => {
let res = await API.get(`/api/token/${tokenId}`);
const { success, message, data } = res.data;
@ -105,13 +110,17 @@ const EditToken = () => {
<Form.Input
label='剩余次数'
name='remain_times'
placeholder={'请输入剩余次数，-1 表示无限制'}
placeholder={'请输入剩余次数'}
onChange={handleInputChange}
value={remain_times}
autoComplete='off'
type='number'
disabled={unlimited_times}
/>
</Form.Field>
<Button type={'button'} onClick={() => {
setUnlimitedTimes();
}}>{unlimited_times ? "取消无限次" : "设置为无限次"}</Button>
<Form.Field>
<Form.Input
label='过期时间'