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

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 {

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") {

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)

model/token.go

@@ -10,17 +10,19 @@ import (
 )
 
 type Token struct {
-	Id             int    `json:"id"`
+	Id                 int    `json:"id"`
-	UserId         int    `json:"user_id"`
+	UserId             int    `json:"user_id"`
-	Key            string `json:"key" gorm:"type:char(48);uniqueIndex"`
+	Key                string `json:"key" gorm:"type:char(48);uniqueIndex"`
-	Status         int    `json:"status" gorm:"default:1"`
+	Status             int    `json:"status" gorm:"default:1"`
-	Name           string `json:"name" gorm:"index" `
+	Name               string `json:"name" gorm:"index" `
-	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
-	RemainQuota    int    `json:"remain_quota" gorm:"default:0"`
+	RemainQuota        int    `json:"remain_quota" gorm:"default:0"`
-	UnlimitedQuota bool   `json:"unlimited_quota" gorm:"default:false"`
+	UnlimitedQuota     bool   `json:"unlimited_quota" gorm:"default:false"`
-	UsedQuota      int    `json:"used_quota" gorm:"default:0"` // used quota
+	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 {

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 <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) => {

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([]);

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) {
+        setIsEdit(props.editingToken.id !== undefined);
-            loadToken().then(
-                () => {
-                    // console.log(inputs);
-                }
-            );
-        } else {
-            setInputs(originInputs);
-        }
     }, [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 }}
+                            <div style={{marginTop: 20}}>
-                            label='数量'
+                                <Typography.Text>新建数量</Typography.Text>
-                            placeholder={'请选择或输入创建令牌的数量'}
+                            </div>
-                            onChange={(value) => handleTokenCountChange(value)}
+                            <AutoComplete
-                            onSelect={(value) => handleTokenCountChange(value)}
+                                style={{ marginTop: 8 }}
-                            value={tokenCount.toString()}
+                                label='数量'
-                            autoComplete='off'
+                                placeholder={'请选择或输入创建令牌的数量'}
-                            type='number'
+                                onChange={(value) => handleTokenCountChange(value)}
-                            data={[
+                                onSelect={(value) => handleTokenCountChange(value)}
-                                { value: 10, label: '10个' },
+                                value={tokenCount.toString()}
-                                { value: 20, label: '20个' },
+                                autoComplete='off'
-                                { value: 30, label: '30个' },
+                                type='number'
-                                { value: 100, label: '100个' },
+                                data={[
-                            ]}
+                                    { value: 10, label: '10个' },
-                            disabled={unlimited_quota}
+                                    { 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>
         </>