feat: 令牌分组

2025-11-04 05:13:41 +08:00 · 2024-09-18 05:19:10 +08:00
parent 5f3798053f
commit 052bc2075b
13 changed files with 248 additions and 120 deletions
--- a/common/user_groups.go
+++ b/common/user_groups.go
@@ -0,0 +1,23 @@
+package common
+
+import (
+	"encoding/json"
+)
+
+var UserUsableGroups = map[string]string{
+	"default": "默认分组",
+	"vip":     "vip分组",
+}
+
+func UserUsableGroups2JSONString() string {
+	jsonBytes, err := json.Marshal(UserUsableGroups)
+	if err != nil {
+		SysError("error marshalling user groups: " + err.Error())
+	}
+	return string(jsonBytes)
+}
+
+func UpdateUserUsableGroupsByJSONString(jsonStr string) error {
+	UserUsableGroups = make(map[string]string)
+	return json.Unmarshal([]byte(jsonStr), &UserUsableGroups)
+}
--- a/controller/group.go
+++ b/controller/group.go
@@ -17,3 +17,18 @@ func GetGroups(c *gin.Context) {
 		"data":    groupNames,
 	})
 }
+
+func GetUserGroups(c *gin.Context) {
+	usableGroups := make(map[string]string)
+	for groupName, _ := range common.GroupRatio {
+		// UserUsableGroups contains the groups that the user can use
+		if _, ok := common.UserUsableGroups[groupName]; ok {
+			usableGroups[groupName] = common.UserUsableGroups[groupName]
+		}
+	}
+	c.JSON(http.StatusOK, gin.H{
+		"success": true,
+		"message": "",
+		"data":    usableGroups,
+	})
+}
--- a/controller/token.go
+++ b/controller/token.go
@@ -135,6 +135,7 @@ func AddToken(c *gin.Context) {
 		ModelLimitsEnabled: token.ModelLimitsEnabled,
 		ModelLimits:        token.ModelLimits,
 		AllowIps:           token.AllowIps,
+		Group:              token.Group,
 	}
 	err = cleanToken.Insert()
 	if err != nil {
@@ -223,6 +224,7 @@ func UpdateToken(c *gin.Context) {
 		cleanToken.ModelLimitsEnabled = token.ModelLimitsEnabled
 		cleanToken.ModelLimits = token.ModelLimits
 		cleanToken.AllowIps = token.AllowIps
+		cleanToken.Group = token.Group
 	}
 	err = cleanToken.Update()
 	if err != nil {
--- a/middleware/auth.go
+++ b/middleware/auth.go
@@ -176,6 +176,7 @@ func TokenAuth() func(c *gin.Context) {
 			c.Set("token_model_limit_enabled", false)
 		}
 		c.Set("allow_ips", token.GetIpLimitsMap())
+		c.Set("token_group", token.Group)
 		if len(parts) > 1 {
 			if model.IsAdmin(token.UserId) {
 				c.Set("specific_channel_id", parts[1])
--- a/middleware/distributor.go
+++ b/middleware/distributor.go
@@ -39,6 +39,15 @@ func Distribute() func(c *gin.Context) {
 			return
 		}
 		userGroup, _ := model.CacheGetUserGroup(userId)
+		tokenGroup := c.GetString("token_group")
+		if tokenGroup != "" {
+			// check group in common.GroupRatio
+			if _, ok := common.GroupRatio[tokenGroup]; !ok {
+				abortWithOpenAiMessage(c, http.StatusForbidden, fmt.Sprintf("分组 %s 已被禁用", tokenGroup))
+				return
+			}
+			userGroup = tokenGroup
+		}
 		c.Set("group", userGroup)
 		if ok {
 			id, err := strconv.Atoi(channelId.(string))
--- a/model/option.go
+++ b/model/option.go
@@ -86,6 +86,7 @@ func InitOptionMap() {
 	common.OptionMap["ModelRatio"] = common.ModelRatio2JSONString()
 	common.OptionMap["ModelPrice"] = common.ModelPrice2JSONString()
 	common.OptionMap["GroupRatio"] = common.GroupRatio2JSONString()
+	common.OptionMap["UserUsableGroups"] = common.UserUsableGroups2JSONString()
 	common.OptionMap["CompletionRatio"] = common.CompletionRatio2JSONString()
 	common.OptionMap["TopUpLink"] = common.TopUpLink
 	common.OptionMap["ChatLink"] = common.ChatLink
@@ -303,6 +304,8 @@ func updateOptionMap(key string, value string) (err error) {
 		err = common.UpdateModelRatioByJSONString(value)
 	case "GroupRatio":
 		err = common.UpdateGroupRatioByJSONString(value)
+	case "UserUsableGroups":
+		err = common.UpdateUserUsableGroupsByJSONString(value)
 	case "CompletionRatio":
 		err = common.UpdateCompletionRatioByJSONString(value)
 	case "ModelPrice":
--- a/model/token.go
+++ b/model/token.go
@@ -25,6 +25,7 @@ type Token struct {
 	ModelLimits        string         `json:"model_limits" gorm:"type:varchar(1024);default:''"`
 	AllowIps           *string        `json:"allow_ips" gorm:"default:''"`
 	UsedQuota          int            `json:"used_quota" gorm:"default:0"` // used quota
+	Group              string         `json:"group" gorm:"default:''"`
 	DeletedAt          gorm.DeletedAt `gorm:"index"`
 }

@@ -153,7 +154,8 @@ 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", "model_limits_enabled", "model_limits", "allow_ips").Updates(token).Error
+	err = DB.Model(token).Select("name", "status", "expired_time", "remain_quota", "unlimited_quota",
+		"model_limits_enabled", "model_limits", "allow_ips", "group").Updates(token).Error
 	return err
 }

--- a/router/api-router.go
+++ b/router/api-router.go
@@ -39,6 +39,7 @@ func SetApiRouter(router *gin.Engine) {
 			//userRoute.POST("/tokenlog", middleware.CriticalRateLimit(), controller.TokenLog)
 			userRoute.GET("/logout", controller.Logout)
 			userRoute.GET("/epay/notify", controller.EpayNotify)
+			userRoute.GET("/groups", controller.GetUserGroups)

 			selfRoute := userRoute.Group("/")
 			selfRoute.Use(middleware.UserAuth())
--- a/web/src/components/OperationSetting.js
+++ b/web/src/components/OperationSetting.js
@@ -23,6 +23,7 @@ const OperationSetting = () => {
    CompletionRatio: '',
    ModelPrice: '',
    GroupRatio: '',
+    UserUsableGroups: '',
    TopUpLink: '',
    ChatLink: '',
    ChatLink2: '', // 添加的新状态变量
@@ -62,6 +63,7 @@ const OperationSetting = () => {
        if (
          item.key === 'ModelRatio' ||
          item.key === 'GroupRatio' ||
+          item.key === 'UserUsableGroups' ||
          item.key === 'CompletionRatio' ||
          item.key === 'ModelPrice'
        ) {
--- a/web/src/components/TokensTable.js
+++ b/web/src/components/TokensTable.js
@@ -8,14 +8,14 @@ import {
 } from '../helpers';

 import { ITEMS_PER_PAGE } from '../constants';
-import { renderQuota } from '../helpers/render';
+import {renderGroup, renderQuota} from '../helpers/render';
 import {
  Button,
  Dropdown,
  Form,
  Modal,
  Popconfirm,
-  Popover,
+  Popover, Space,
  SplitButtonGroup,
  Table,
  Tag,
@@ -119,7 +119,12 @@ const TokensTable = () => {
      dataIndex: 'status',
      key: 'status',
      render: (text, record, index) => {
-        return <div>{renderStatus(text, record.model_limits_enabled)}</div>;
+        return <div>
+          <Space>
+            {renderStatus(text, record.model_limits_enabled)}
+            {renderGroup(record.group)}
+          </Space>
+        </div>;
      },
    },
    {
--- a/web/src/helpers/render.js
+++ b/web/src/helpers/render.js
@@ -15,8 +15,8 @@ export function renderText(text, limit) {
 export function renderGroup(group) {
  if (group === '') {
    return (
-      <Tag size='large' key='default'>
-        unknown
+      <Tag size='large' key='default' color={stringToColor('default')}>
+        default
      </Tag>
    );
  }
--- a/web/src/pages/Setting/Operation/SettingsMagnification.js
+++ b/web/src/pages/Setting/Operation/SettingsMagnification.js
@@ -16,7 +16,8 @@ export default function SettingsMagnification(props) {
    ModelPrice: '',
    ModelRatio: '',
    CompletionRatio: '',
-    GroupRatio: ''
+    GroupRatio: '',
+    UserUsableGroups: ''
  });
  const refForm = useRef();
  const [inputsRow, setInputsRow] = useState(inputs);
@@ -213,6 +214,33 @@ export default function SettingsMagnification(props) {
              />
            </Col>
          </Row>
+          <Row gutter={16}>
+            <Col span={16}>
+              <Form.TextArea
+                  label={'用户可选分组'}
+                  extraText={''}
+                  placeholder={'为一个 JSON 文本，键为分组名称，值为倍率'}
+                  field={'UserUsableGroups'}
+                  autosize={{ minRows: 6, maxRows: 12 }}
+                  trigger='blur'
+                  stopValidateWithError
+                  rules={[
+                    {
+                      validator: (rule, value) => {
+                        return verifyJSON(value);
+                      },
+                      message: '不是合法的 JSON 字符串'
+                    }
+                  ]}
+                  onChange={(value) =>
+                      setInputs({
+                        ...inputs,
+                        UserUsableGroups: value
+                      })
+                  }
+              />
+            </Col>
+          </Row>
        </Form.Section>
      </Form>
      <Space>
--- a/web/src/pages/Token/EditToken.js
+++ b/web/src/pages/Token/EditToken.js
@@ -35,6 +35,7 @@ const EditToken = (props) => {
    model_limits_enabled: false,
    model_limits: [],
    allow_ips: '',
+    group: '',
  };
  const [inputs, setInputs] = useState(originInputs);
  const {
@@ -44,10 +45,12 @@ const EditToken = (props) => {
    unlimited_quota,
    model_limits_enabled,
    model_limits,
-    allow_ips
+    allow_ips,
+    group
  } = inputs;
  // const [visible, setVisible] = useState(false);
  const [models, setModels] = useState({});
+  const [groups, setGroups] = useState([]);
  const navigate = useNavigate();
  const handleInputChange = (name, value) => {
    setInputs((inputs) => ({ ...inputs, [name]: value }));
@@ -88,6 +91,22 @@ const EditToken = (props) => {
    }
  };

+  const loadGroups = async () => {
+    let res = await API.get(`/api/user/groups`);
+    const { success, message, data } = res.data;
+    if (success) {
+      // return data is a map, key is group name, value is group description
+      // label is group description, value is group name
+        let localGroupOptions = Object.keys(data).map((group) => ({
+            label: data[group],
+            value: group,
+        }));
+        setGroups(localGroupOptions);
+    } else {
+      showError(message);
+    }
+  };
+
  const loadToken = async () => {
    setLoading(true);
    let res = await API.get(`/api/token/${props.editingToken.id}`);
@@ -120,6 +139,7 @@ const EditToken = (props) => {
      });
    }
    loadModels();
+    loadGroups();
  }, [isEdit]);

  // 新增 state 变量 tokenCount 来记录用户想要创建的令牌数量，默认为 1
@@ -253,150 +273,150 @@ const EditToken = (props) => {
      >
        <Spin spinning={loading}>
          <Input
-            style={{ marginTop: 20 }}
-            label='名称'
-            name='name'
-            placeholder={'请输入名称'}
-            onChange={(value) => handleInputChange('name', value)}
-            value={name}
-            autoComplete='new-password'
-            required={!isEdit}
+              style={{marginTop: 20}}
+              label='名称'
+              name='name'
+              placeholder={'请输入名称'}
+              onChange={(value) => handleInputChange('name', value)}
+              value={name}
+              autoComplete='new-password'
+              required={!isEdit}
          />
-          <Divider />
+          <Divider/>
          <DatePicker
-            label='过期时间'
-            name='expired_time'
-            placeholder={'请选择过期时间'}
-            onChange={(value) => handleInputChange('expired_time', value)}
-            value={expired_time}
-            autoComplete='new-password'
-            type='dateTime'
+              label='过期时间'
+              name='expired_time'
+              placeholder={'请选择过期时间'}
+              onChange={(value) => handleInputChange('expired_time', value)}
+              value={expired_time}
+              autoComplete='new-password'
+              type='dateTime'
          />
-          <div style={{ marginTop: 20 }}>
+          <div style={{marginTop: 20}}>
            <Space>
              <Button
-                type={'tertiary'}
-                onClick={() => {
-                  setExpiredTime(0, 0, 0, 0);
-                }}
+                  type={'tertiary'}
+                  onClick={() => {
+                    setExpiredTime(0, 0, 0, 0);
+                  }}
              >
                永不过期
              </Button>
              <Button
-                type={'tertiary'}
-                onClick={() => {
-                  setExpiredTime(0, 0, 1, 0);
-                }}
+                  type={'tertiary'}
+                  onClick={() => {
+                    setExpiredTime(0, 0, 1, 0);
+                  }}
              >
                一小时
              </Button>
              <Button
-                type={'tertiary'}
-                onClick={() => {
-                  setExpiredTime(1, 0, 0, 0);
-                }}
+                  type={'tertiary'}
+                  onClick={() => {
+                    setExpiredTime(1, 0, 0, 0);
+                  }}
              >
                一个月
              </Button>
              <Button
-                type={'tertiary'}
-                onClick={() => {
-                  setExpiredTime(0, 1, 0, 0);
-                }}
+                  type={'tertiary'}
+                  onClick={() => {
+                    setExpiredTime(0, 1, 0, 0);
+                  }}
              >
                一天
              </Button>
            </Space>
          </div>

-          <Divider />
+          <Divider/>
          <Banner
-            type={'warning'}
-            description={
-              '注意，令牌的额度仅用于限制令牌本身的最大额度使用量，实际的使用受到账户的剩余额度限制。'
-            }
+              type={'warning'}
+              description={
+                '注意，令牌的额度仅用于限制令牌本身的最大额度使用量，实际的使用受到账户的剩余额度限制。'
+              }
          ></Banner>
-          <div style={{ marginTop: 20 }}>
+          <div style={{marginTop: 20}}>
            <Typography.Text>{`额度${renderQuotaWithPrompt(remain_quota)}`}</Typography.Text>
          </div>
          <AutoComplete
-            style={{ marginTop: 8 }}
-            name='remain_quota'
-            placeholder={'请输入额度'}
-            onChange={(value) => handleInputChange('remain_quota', value)}
-            value={remain_quota}
-            autoComplete='new-password'
-            type='number'
-            // position={'top'}
-            data={[
-              { value: 500000, label: '1$' },
-              { value: 5000000, label: '10$' },
-              { value: 25000000, label: '50$' },
-              { value: 50000000, label: '100$' },
-              { value: 250000000, label: '500$' },
-              { value: 500000000, label: '1000$' },
-            ]}
-            disabled={unlimited_quota}
+              style={{marginTop: 8}}
+              name='remain_quota'
+              placeholder={'请输入额度'}
+              onChange={(value) => handleInputChange('remain_quota', value)}
+              value={remain_quota}
+              autoComplete='new-password'
+              type='number'
+              // position={'top'}
+              data={[
+                {value: 500000, label: '1$'},
+                {value: 5000000, label: '10$'},
+                {value: 25000000, label: '50$'},
+                {value: 50000000, label: '100$'},
+                {value: 250000000, label: '500$'},
+                {value: 500000000, label: '1000$'},
+              ]}
+              disabled={unlimited_quota}
          />

          {!isEdit && (
-            <>
-              <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 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>
            <Button
-              style={{ marginTop: 8 }}
-              type={'warning'}
-              onClick={() => {
-                setUnlimitedQuota();
-              }}
+                style={{marginTop: 8}}
+                type={'warning'}
+                onClick={() => {
+                  setUnlimitedQuota();
+                }}
            >
              {unlimited_quota ? '取消无限额度' : '设为无限额度'}
            </Button>
          </div>
-          <Divider />
-          <div style={{ marginTop: 10 }}>
+          <Divider/>
+          <div style={{marginTop: 10}}>
            <Typography.Text>IP白名单（请勿过度信任此功能）</Typography.Text>
          </div>
          <TextArea
-            label='IP白名单'
-            name='allow_ips'
-            placeholder={'允许的IP，一行一个'}
-            onChange={(value) => {
-              handleInputChange('allow_ips', value);
-            }}
-            value={inputs.allow_ips}
-            style={{ fontFamily: 'JetBrains Mono, Consolas' }}
+              label='IP白名单'
+              name='allow_ips'
+              placeholder={'允许的IP，一行一个'}
+              onChange={(value) => {
+                handleInputChange('allow_ips', value);
+              }}
+              value={inputs.allow_ips}
+              style={{fontFamily: 'JetBrains Mono, Consolas'}}
          />
-          <div style={{ marginTop: 10, display: 'flex' }}>
+          <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)
-                }
+                  name='model_limits_enabled'
+                  checked={model_limits_enabled}
+                  onChange={(e) =>
+                      handleInputChange('model_limits_enabled', e.target.checked)
+                  }
              ></Checkbox>
              <Typography.Text>
                启用模型限制（非必要，不建议启用）
@@ -405,19 +425,36 @@ const EditToken = (props) => {
          </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}
+              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}
+          />
+
+          <div style={{marginTop: 10}}>
+            <Typography.Text>令牌分组，不选为默认分组</Typography.Text>
+          </div>
+          <Select
+              style={{marginTop: 8}}
+              placeholder={'令牌分组，不选为默认分组'}
+              name='gruop'
+              required
+              selection
+              onChange={(value) => {
+                handleInputChange('group', value);
+              }}
+              value={inputs.group}
+              autoComplete='new-password'
+              optionList={groups}
          />
        </Spin>
      </SideSheet>