diff --git a/controller/token.go b/controller/token.go index 0f37fd0..157cb2f 100644 --- a/controller/token.go +++ b/controller/token.go @@ -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 { diff --git a/middleware/auth.go b/middleware/auth.go index c0ff074..e12b81b 100644 --- a/middleware/auth.go +++ b/middleware/auth.go @@ -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") { diff --git a/middleware/distributor.go b/middleware/distributor.go index 6d64e30..a70ed41 100644 --- a/middleware/distributor.go +++ b/middleware/distributor.go @@ -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) diff --git a/model/token.go b/model/token.go index b5f9230..3c6f35d 100644 --- a/model/token.go +++ b/model/token.go @@ -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 { diff --git a/web/src/components/TokensTable.js b/web/src/components/TokensTable.js index aa8fd57..5b685c0 100644 --- a/web/src/components/TokensTable.js +++ b/web/src/components/TokensTable.js @@ -43,10 +43,14 @@ function renderTimestamp(timestamp) { ); } -function renderStatus(status) { +function renderStatus(status, model_limits_enabled = false) { switch (status) { case 1: - return 已启用; + if (model_limits_enabled) { + return 已启用：限制模型; + } else { + return 已启用; + } case 2: return 已禁用 ; case 3: @@ -78,7 +82,7 @@ const TokensTable = () => { render: (text, record, index) => { return (
- {renderStatus(text)} + {renderStatus(text, record.model_limits_enabled)}
); }, @@ -224,6 +228,11 @@ const TokensTable = () => { const closeEdit = () => { setShowEdit(false); + setTimeout(() => { + setEditingToken({ + id: undefined, + }); + }, 500); } const setTokensFormat = (tokens) => { diff --git a/web/src/pages/Detail/index.js b/web/src/pages/Detail/index.js index c36e54d..7f582da 100644 --- a/web/src/pages/Detail/index.js +++ b/web/src/pages/Detail/index.js @@ -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([]); diff --git a/web/src/pages/Token/EditToken.js b/web/src/pages/Token/EditToken.js index 0784c7a..d7fd909 100644 --- a/web/src/pages/Token/EditToken.js +++ b/web/src/pages/Token/EditToken.js @@ -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} /> -
- 新建数量 -
+ {!isEdit && ( - 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} - /> + <> +
+ 新建数量 +
+ 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} + /> + )}
@@ -273,6 +315,34 @@ const EditToken = (props) => { setUnlimitedQuota(); }}>{unlimited_quota ? '取消无限额度' : '设为无限额度'}
+ +
+ + handleInputChange('model_limits_enabled', e.target.checked)} + > + + 启用模型限制（非必要，不建议启用） + +
+ +