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

feat: 令牌分组

Browse Source
This commit is contained in:
CalciumIon 2024-09-18 05:19:10 +08:00
parent 5f3798053f
commit 052bc2075b
13 changed files with 248 additions and 120 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

23
common/user_groups.go Normal file
View File

@ -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)
}

15
controller/group.go
View File

@ -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,
})
}

2
controller/token.go
View File

@ -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 {

1
middleware/auth.go
View File

@ -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])

9
middleware/distributor.go
View File

@ -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))

3
model/option.go
View File

@ -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":

4
model/token.go
View File

@ -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
}

1
router/api-router.go
View File

@ -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())

2
web/src/components/OperationSetting.js
View File

@ -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'
) {

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

@ -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>;
},
},
{

4
web/src/helpers/render.js
View File

@ -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>
);
}

30
web/src/pages/Setting/Operation/SettingsMagnification.js
View File

@ -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>

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

@ -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
@ -419,6 +439,23 @@ const EditToken = (props) => {
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>
</>