xinyin025/new-api
1
0
Fork 0
mirror of https://github.com/linux-do/new-api.git synced 2025-09-20 01:16:36 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

feat: 可设置令牌能调用的模型

Browse Source
This commit is contained in:
CaIon 2024-01-08 16:23:54 +08:00
parent 8f36a995ef
commit 1244963e81
7 changed files with 192 additions and 52 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
controller/token.go
View File

@ -217,6 +217,8 @@ func UpdateToken(c *gin.Context) {
cleanToken.ExpiredTime = token.ExpiredTime
cleanToken.RemainQuota = token.RemainQuota
cleanToken.UnlimitedQuota = token.UnlimitedQuota
cleanToken.ModelLimitsEnabled = token.ModelLimitsEnabled
cleanToken.ModelLimits = token.ModelLimits
}
err = cleanToken.Update()
if err != nil {

6
middleware/auth.go
View File

@ -115,6 +115,12 @@ func TokenAuth() func(c *gin.Context) {
c.Set("id", token.UserId)
c.Set("token_id", token.Id)
c.Set("token_name", token.Name)
if token.ModelLimitsEnabled {
c.Set("token_model_limit_enabled", true)
c.Set("token_model_limit", token.GetModelLimitsMap())
} else {
c.Set("token_model_limit_enabled", false)
}
requestURL := c.Request.URL.String()
consumeQuota := true
if strings.HasPrefix(requestURL, "/v1/models") {

21
middleware/distributor.go
View File

@ -77,6 +77,27 @@ func Distribute() func(c *gin.Context) {
}
}
}
// check token model mapping
modelLimitEnable := c.GetBool("token_model_limit_enabled")
if modelLimitEnable {
s, ok := c.Get("token_model_limit")
var tokenModelLimit map[string]bool
if ok {
tokenModelLimit = s.(map[string]bool)
} else {
tokenModelLimit = map[string]bool{}
}
if tokenModelLimit != nil {
if _, ok := tokenModelLimit[modelRequest.Model]; !ok {
abortWithMessage(c, http.StatusForbidden, "该令牌无权访问模型 "+modelRequest.Model)
return
}
} else {
// token model limit is empty, all models are not allowed
abortWithMessage(c, http.StatusForbidden, "该令牌无权访问任何模型")
return
}
}
channel, err = model.CacheGetRandomSatisfiedChannel(userGroup, modelRequest.Model)
if err != nil {
message := fmt.Sprintf("当前分组 %s 下对于模型 %s 无可用渠道", userGroup, modelRequest.Model)

56
model/token.go
View File

@ -10,17 +10,19 @@ import (
)
type Token struct {
Id int `json:"id"`
UserId int `json:"user_id"`
Key string `json:"key" gorm:"type:char(48);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
RemainQuota int `json:"remain_quota" gorm:"default:0"`
UnlimitedQuota bool `json:"unlimited_quota" gorm:"default:false"`
UsedQuota int `json:"used_quota" gorm:"default:0"` // used quota
Id int `json:"id"`
UserId int `json:"user_id"`
Key string `json:"key" gorm:"type:char(48);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
RemainQuota int `json:"remain_quota" gorm:"default:0"`
UnlimitedQuota bool `json:"unlimited_quota" gorm:"default:false"`
ModelLimitsEnabled bool `json:"model_limits_enabled" gorm:"default:false"`
ModelLimits string `json:"model_limits" gorm:"type:varchar(1024);default:''"`
UsedQuota int `json:"used_quota" gorm:"default:0"` // used quota
}
func GetAllUserTokens(userId int, startIdx int, num int) ([]*Token, error) {
@ -107,7 +109,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_quota", "unlimited_quota").Updates(token).Error
err = DB.Model(token).Select("name", "status", "expired_time", "remain_quota", "unlimited_quota", "model_limits_enabled", "model_limits").Updates(token).Error
return err
}
@ -122,6 +124,36 @@ func (token *Token) Delete() error {
return err
}
func (token *Token) IsModelLimitsEnabled() bool {
return token.ModelLimitsEnabled
}
func (token *Token) GetModelLimits() []string {
if token.ModelLimits == "" {
return []string{}
}
return strings.Split(token.ModelLimits, ",")
}
func (token *Token) GetModelLimitsMap() map[string]bool {
limits := token.GetModelLimits()
limitsMap := make(map[string]bool)
for _, limit := range limits {
limitsMap[limit] = true
}
return limitsMap
}
func DisableModelLimits(tokenId int) error {
token, err := GetTokenById(tokenId)
if err != nil {
return err
}
token.ModelLimitsEnabled = false
token.ModelLimits = ""
return token.Update()
}
func DeleteTokenById(id int, userId int) (err error) {
// Why we need userId here? In case user want to delete other's token.
if id == 0 || userId == 0 {

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

@ -43,10 +43,14 @@ function renderTimestamp(timestamp) {
);
}
function renderStatus(status) {
function renderStatus(status, model_limits_enabled = false) {
switch (status) {
case 1:
return <Tag color='green' size='large'>已启用</Tag>;
if (model_limits_enabled) {
return <Tag color='green' size='large'>已启用：限制模型</Tag>;
} else {
return <Tag color='green' size='large'>已启用</Tag>;
}
case 2:
return <Tag color='red' size='large'> 已禁用 </Tag>;
case 3:
@ -78,7 +82,7 @@ const TokensTable = () => {
render: (text, record, index) => {
return (
<div>
{renderStatus(text)}
{renderStatus(text, record.model_limits_enabled)}
</div>
);
},
@ -224,6 +228,11 @@ const TokensTable = () => {
const closeEdit = () => {
setShowEdit(false);
setTimeout(() => {
setEditingToken({
id: undefined,
});
}, 500);
}
const setTokensFormat = (tokens) => {

2
web/src/pages/Detail/index.js
View File

@ -21,7 +21,7 @@ const Detail = (props) => {
const initialized = useRef(false)
const [modelDataChart, setModelDataChart] = useState(null);
const [modelDataPieChart, setModelDataPieChart] = useState(null);
const [loading, setLoading] = useState(true);
const [loading, setLoading] = useState(false);
const [quotaData, setQuotaData] = useState([]);
const [quotaDataPie, setQuotaDataPie] = useState([]);
const [quotaDataLine, setQuotaDataLine] = useState([]);

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

@ -2,22 +2,37 @@ import React, {useEffect, useRef, useState} from 'react';
import {useParams, useNavigate} from 'react-router-dom';
import {API, isMobile, showError, showSuccess, timestamp2string} from '../../helpers';
import {renderQuota, renderQuotaWithPrompt} from '../../helpers/render';
import {Layout, SideSheet, Button, Space, Spin, Banner, Input, DatePicker, AutoComplete, Typography} from "@douyinfe/semi-ui";
import {
Layout,
SideSheet,
Button,
Space,
Spin,
Banner,
Input,
DatePicker,
AutoComplete,
Typography,
Checkbox, Select
} from "@douyinfe/semi-ui";
import Title from "@douyinfe/semi-ui/lib/es/typography/title";
import {Divider} from "semantic-ui-react";
const EditToken = (props) => {
const isEdit = props.editingToken.id !== undefined;
const [isEdit, setIsEdit] = useState(false);
const [loading, setLoading] = useState(isEdit);
const originInputs = {
name: '',
remain_quota: isEdit ? 0 : 500000,
expired_time: -1,
unlimited_quota: false
unlimited_quota: false,
model_limits_enabled: false,
model_limits: [],
};
const [inputs, setInputs] = useState(originInputs);
const {name, remain_quota, expired_time, unlimited_quota} = inputs;
const {name, remain_quota, expired_time, unlimited_quota, model_limits_enabled, model_limits} = inputs;
// const [visible, setVisible] = useState(false);
const [models, setModels] = useState({});
const navigate = useNavigate();
const handleInputChange = (name, value) => {
setInputs((inputs) => ({...inputs, [name]: value}));
@ -44,6 +59,20 @@ const EditToken = (props) => {
setInputs({...inputs, unlimited_quota: !unlimited_quota});
};
const loadModels = async () => {
let res = await API.get(`/api/user/models`);
const {success, message, data} = res.data;
if (success) {
let localModelOptions = data.map((model) => ({
label: model,
value: model
}));
setModels(localModelOptions);
} else {
showError(message);
}
}
const loadToken = async () => {
setLoading(true);
let res = await API.get(`/api/token/${props.editingToken.id}`);
@ -52,6 +81,11 @@ const EditToken = (props) => {
if (data.expired_time !== -1) {
data.expired_time = timestamp2string(data.expired_time);
}
if (data.model_limits !== '') {
data.model_limits = data.model_limits.split(',');
} else {
data.model_limits = [];
}
setInputs(data);
} else {
showError(message);
@ -59,17 +93,22 @@ const EditToken = (props) => {
setLoading(false);
};
useEffect(() => {
if (isEdit) {
loadToken().then(
() => {
// console.log(inputs);
}
);
} else {
setInputs(originInputs);
}
setIsEdit(props.editingToken.id !== undefined);
}, [props.editingToken.id]);
useEffect(() => {
if (!isEdit) {
setInputs(originInputs);
} else {
loadToken().then(
() => {
// console.log(inputs);
}
);
}
loadModels();
}, [isEdit]);
// 新增 state 变量 tokenCount 来记录用户想要创建的令牌数量，默认为 1
const [tokenCount, setTokenCount] = useState(1);
@ -107,7 +146,7 @@ const EditToken = (props) => {
}
localInputs.expired_time = Math.ceil(time / 1000);
}
localInputs.model_limits = localInputs.model_limits.join(',');
let res = await API.put(`/api/token/`, {...localInputs, id: parseInt(props.editingToken.id)});
const {success, message} = res.data;
if (success) {
@ -137,7 +176,7 @@ const EditToken = (props) => {
}
localInputs.expired_time = Math.ceil(time / 1000);
}
localInputs.model_limits = localInputs.model_limits.join(',');
let res = await API.post(`/api/token/`, localInputs);
const {success, message} = res.data;
@ -234,7 +273,7 @@ const EditToken = (props) => {
value={remain_quota}
autoComplete='new-password'
type='number'
position={'top'}
// position={'top'}
data={[
{value: 500000, label: '1$'},
{value: 5000000, label: '10$'},
@ -245,27 +284,30 @@ const EditToken = (props) => {
]}
disabled={unlimited_quota}
/>
<div style={{marginTop: 20}}>
<Typography.Text>新建数量</Typography.Text>
</div>
{!isEdit && (
<AutoComplete
style={{ marginTop: 8 }}
label='数量'
placeholder={'请选择或输入创建令牌的数量'}
onChange={(value) => handleTokenCountChange(value)}
onSelect={(value) => handleTokenCountChange(value)}
value={tokenCount.toString()}
autoComplete='off'
type='number'
data={[
{ value: 10, label: '10个' },
{ value: 20, label: '20个' },
{ value: 30, label: '30个' },
{ value: 100, label: '100个' },
]}
disabled={unlimited_quota}
/>
<>
<div style={{marginTop: 20}}>
<Typography.Text>新建数量</Typography.Text>
</div>
<AutoComplete
style={{ marginTop: 8 }}
label='数量'
placeholder={'请选择或输入创建令牌的数量'}
onChange={(value) => handleTokenCountChange(value)}
onSelect={(value) => handleTokenCountChange(value)}
value={tokenCount.toString()}
autoComplete='off'
type='number'
data={[
{ value: 10, label: '10个' },
{ value: 20, label: '20个' },
{ value: 30, label: '30个' },
{ value: 100, label: '100个' },
]}
disabled={unlimited_quota}
/>
</>
)}
<div>
@ -273,6 +315,34 @@ const EditToken = (props) => {
setUnlimitedQuota();
}}>{unlimited_quota ? '取消无限额度' : '设为无限额度'}</Button>
</div>
<Divider/>
<div style={{marginTop: 10, display: 'flex'}}>
<Space>
<Checkbox
name='model_limits_enabled'
checked={model_limits_enabled}
onChange={(e) => handleInputChange('model_limits_enabled', e.target.checked)}
>
</Checkbox>
<Typography.Text>启用模型限制（非必要，不建议启用）</Typography.Text>
</Space>
</div>
<Select
style={{marginTop: 8}}
placeholder={'请选择该渠道所支持的模型'}
name='models'
required
multiple
selection
onChange={value => {
handleInputChange('model_limits', value);
}}
value={inputs.model_limits}
autoComplete='new-password'
optionList={models}
disabled={!model_limits_enabled}
/>
</Spin>
</SideSheet>
</>