fix: only reduce remain times when request /v1/chat/completions
(close #15)
BREAKING CHANGE: now remain_times is -1 doesn't mean unlimited times anymore!
This commit is contained in:
parent
eb8f43acb5
commit
109736cc05
@ -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)
|
||||
|
@ -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{
|
||||
|
@ -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])
|
||||
}
|
||||
|
@ -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
|
||||
}
|
||||
|
@ -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>
|
||||
|
@ -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='过期时间'
|
||||
|
Loading…
Reference in New Issue
Block a user