feat: pricing page support multi groups #487

common/model-ratio.go

@@ -375,6 +375,9 @@ func GetCompletionRatio(name string) float64 {
 		return 3
 	}
 	if strings.HasPrefix(name, "gemini-") {
+		if strings.Contains(name, "flash") {
+			return 4
+		}
 		return 3
 	}
 	if strings.HasPrefix(name, "command") {

controller/model.go

@@ -137,15 +137,6 @@ func init() {
 }
 
 func ListModels(c *gin.Context) {
-	userId := c.GetInt("id")
-	user, err := model.GetUserById(userId, true)
-	if err != nil {
-		c.JSON(http.StatusOK, gin.H{
-			"success": false,
-			"message": err.Error(),
-		})
-		return
-	}
 	userOpenAiModels := make([]dto.OpenAIModels, 0)
 	permission := getPermission()
 
@@ -174,7 +165,21 @@ func ListModels(c *gin.Context) {
 			}
 		}
 	} else {
-		models := model.GetGroupModels(user.Group)
+		userId := c.GetInt("id")
+		userGroup, err := model.GetUserGroup(userId)
+		if err != nil {
+			c.JSON(http.StatusOK, gin.H{
+				"success": false,
+				"message": "get user group failed",
+			})
+			return
+		}
+		group := userGroup
+		tokenGroup := c.GetString("token_group")
+		if tokenGroup != "" {
+			group = tokenGroup
+		}
+		models := model.GetGroupModels(group)
 		for _, s := range models {
 			if _, ok := openAIModelsMap[s]; ok {
 				userOpenAiModels = append(userOpenAiModels, openAIModelsMap[s])

controller/pricing.go

@@ -7,18 +7,11 @@ import (
 )
 
 func GetPricing(c *gin.Context) {
-	userId := c.GetInt("id")
-	// if no login, get default group ratio
-	groupRatio := common.GetGroupRatio("default")
-	group, err := model.CacheGetUserGroup(userId)
-	if err == nil {
-		groupRatio = common.GetGroupRatio(group)
-	}
-	pricing := model.GetPricing(group)
+	pricing := model.GetPricing()
 	c.JSON(200, gin.H{
 		"success":     true,
 		"data":        pricing,
-		"group_ratio": groupRatio,
+		"group_ratio": common.GroupRatio,
 	})
 }
 

middleware/distributor.go

@@ -41,9 +41,14 @@ func Distribute() func(c *gin.Context) {
 		userGroup, _ := model.CacheGetUserGroup(userId)
 		tokenGroup := c.GetString("token_group")
 		if tokenGroup != "" {
+			// check common.UserUsableGroups[userGroup]
+			if _, ok := common.UserUsableGroups[tokenGroup]; !ok {
+				abortWithOpenAiMessage(c, http.StatusForbidden, fmt.Sprintf("令牌分组 %s 已被禁用", tokenGroup))
+				return
+			}
 			// check group in common.GroupRatio
 			if _, ok := common.GroupRatio[tokenGroup]; !ok {
-				abortWithOpenAiMessage(c, http.StatusForbidden, fmt.Sprintf("分组 %s 已被禁用", tokenGroup))
+				abortWithOpenAiMessage(c, http.StatusForbidden, fmt.Sprintf("分组 %s 已被弃用", tokenGroup))
 				return
 			}
 			userGroup = tokenGroup

model/ability.go

@@ -36,6 +36,12 @@ func GetEnabledModels() []string {
 	return models
 }
 
+func GetAllEnableAbilities() []Ability {
+	var abilities []Ability
+	DB.Find(&abilities, "enabled = ?", true)
+	return abilities
+}
+
 func getPriority(group string, model string, retry int) (int, error) {
 	groupCol := "`group`"
 	trueVal := "1"

model/pricing.go

@@ -7,14 +7,13 @@ import (
 )
 
 type Pricing struct {
-	Available       bool     `json:"available"`
 	ModelName       string   `json:"model_name"`
 	QuotaType       int      `json:"quota_type"`
 	ModelRatio      float64  `json:"model_ratio"`
 	ModelPrice      float64  `json:"model_price"`
 	OwnerBy         string   `json:"owner_by"`
 	CompletionRatio float64  `json:"completion_ratio"`
-	EnableGroup     []string `json:"enable_group,omitempty"`
+	EnableGroup     []string `json:"enable_groups,omitempty"`
 }
 
 var (
@@ -23,40 +22,47 @@ var (
 	updatePricingLock  sync.Mutex
 )
 
-func GetPricing(group string) []Pricing {
+func GetPricing() []Pricing {
 	updatePricingLock.Lock()
 	defer updatePricingLock.Unlock()
 
 	if time.Since(lastGetPricingTime) > time.Minute*1 || len(pricingMap) == 0 {
 		updatePricing()
 	}
-	if group != "" {
-		userPricingMap := make([]Pricing, 0)
-		models := GetGroupModels(group)
-		for _, pricing := range pricingMap {
-			if !common.StringsContains(models, pricing.ModelName) {
-				pricing.Available = false
-			}
-			userPricingMap = append(userPricingMap, pricing)
-		}
-		return userPricingMap
-	}
+	//if group != "" {
+	//	userPricingMap := make([]Pricing, 0)
+	//	models := GetGroupModels(group)
+	//	for _, pricing := range pricingMap {
+	//		if !common.StringsContains(models, pricing.ModelName) {
+	//			pricing.Available = false
+	//		}
+	//		userPricingMap = append(userPricingMap, pricing)
+	//	}
+	//	return userPricingMap
+	//}
 	return pricingMap
 }
 
 func updatePricing() {
 	//modelRatios := common.GetModelRatios()
-	enabledModels := GetEnabledModels()
-	allModels := make(map[string]int)
-	for i, model := range enabledModels {
-		allModels[model] = i
+	enableAbilities := GetAllEnableAbilities()
+	modelGroupsMap := make(map[string][]string)
+	for _, ability := range enableAbilities {
+		groups := modelGroupsMap[ability.Model]
+		if groups == nil {
+			groups = make([]string, 0)
+		}
+		if !common.StringsContains(groups, ability.Group) {
+			groups = append(groups, ability.Group)
+		}
+		modelGroupsMap[ability.Model] = groups
 	}
 
 	pricingMap = make([]Pricing, 0)
-	for model, _ := range allModels {
+	for model, groups := range modelGroupsMap {
 		pricing := Pricing{
-			Available: true,
-			ModelName: model,
+			ModelName:   model,
+			EnableGroup: groups,
 		}
 		modelPrice, findPrice := common.GetModelPrice(model, false)
 		if findPrice {

web/src/components/HeaderBar.js

@@ -36,7 +36,7 @@ let buttons = [
     text: '首页',
     itemKey: 'home',
     to: '/',
-    icon: <IconHomeStroked />,
+    // icon: <IconHomeStroked />,
   },
   // {
   //   text: '模型价格',

web/src/components/ModelPricing.js

@@ -1,5 +1,5 @@
 import React, { useContext, useEffect, useRef, useMemo, useState } from 'react';
-import { API, copy, showError, showSuccess } from '../helpers';
+import { API, copy, showError, showInfo, showSuccess } from '../helpers';
 
 import {
   Banner,
@@ -87,6 +87,7 @@ const ModelPricing = () => {
   const [selectedRowKeys, setSelectedRowKeys] = useState([]);
   const [modalImageUrl, setModalImageUrl] = useState('');
   const [isModalOpenurl, setIsModalOpenurl] = useState(false);
+  const [selectedGroup, setSelectedGroup] = useState('default');
 
   const rowSelection = useMemo(
       () => ({
@@ -120,7 +121,8 @@ const ModelPricing = () => {
       title: '可用性',
       dataIndex: 'available',
       render: (text, record, index) => {
-        return renderAvailable(text);
+         // if record.enable_groups contains selectedGroup, then available is true
+        return renderAvailable(record.enable_groups.includes(selectedGroup));
       },
       sorter: (a, b) => a.available - b.available,
     },
@@ -166,6 +168,43 @@ const ModelPricing = () => {
       },
       sorter: (a, b) => a.quota_type - b.quota_type,
     },
+    {
+      title: '可用分组',
+      dataIndex: 'enable_groups',
+      render: (text, record, index) => {
+        // enable_groups is a string array
+        return (
+          <Space>
+            {text.map((group) => {
+              if (group === selectedGroup) {
+                return (
+                  <Tag
+                    color='blue'
+                    size='large'
+                    prefixIcon={<IconVerify />}
+                  >
+                    {group}
+                  </Tag>
+                );
+              } else {
+                return (
+                  <Tag
+                    color='blue'
+                    size='large'
+                    onClick={() => {
+                      setSelectedGroup(group);
+                      showInfo('当前查看的分组为：' + group + '，倍率为：' + groupRatio[group]);
+                    }}
+                  >
+                    {group}
+                  </Tag>
+                );
+              }
+            })}
+          </Space>
+        );
+      },
+    },
     {
       title: () => (
         <span style={{'display':'flex','alignItems':'center'}}>
@@ -201,6 +240,8 @@ const ModelPricing = () => {
             <Text>模型：{record.quota_type === 0 ? text : '无'}</Text>
             <br />
             <Text>补全：{record.quota_type === 0 ? completionRatio : '无'}</Text>
+            <br />
+            <Text>分组：{groupRatio[selectedGroup]}</Text>
           </>
         );
         return <div>{content}</div>;
@@ -213,11 +254,11 @@ const ModelPricing = () => {
         let content = text;
         if (record.quota_type === 0) {
           // 这里的 *2 是因为 1倍率=0.002刀，请勿删除
-          let inputRatioPrice = record.model_ratio * 2 * record.group_ratio;
+          let inputRatioPrice = record.model_ratio * 2 * groupRatio[selectedGroup];
           let completionRatioPrice =
             record.model_ratio *
             record.completion_ratio * 2 *
-            record.group_ratio;
+            groupRatio[selectedGroup];
           content = (
             <>
               <Text>提示 ${inputRatioPrice} / 1M tokens</Text>
@@ -226,7 +267,7 @@ const ModelPricing = () => {
             </>
           );
         } else {
-          let price = parseFloat(text) * record.group_ratio;
+          let price = parseFloat(text) * groupRatio[selectedGroup];
           content = <>模型价格：${price}</>;
         }
         return <div>{content}</div>;
@@ -237,12 +278,12 @@ const ModelPricing = () => {
   const [models, setModels] = useState([]);
   const [loading, setLoading] = useState(true);
   const [userState, userDispatch] = useContext(UserContext);
-  const [groupRatio, setGroupRatio] = useState(1);
+  const [groupRatio, setGroupRatio] = useState({});
 
   const setModelsFormat = (models, groupRatio) => {
     for (let i = 0; i < models.length; i++) {
       models[i].key = models[i].model_name;
-      models[i].group_ratio = groupRatio;
+      models[i].group_ratio = groupRatio[models[i].model_name];
     }
     // sort by quota_type
     models.sort((a, b) => {
@@ -275,6 +316,7 @@ const ModelPricing = () => {
     const { success, message, data, group_ratio } = res.data;
     if (success) {
       setGroupRatio(group_ratio);
+      setSelectedGroup(userState.user ? userState.user.group : 'default')
       setModelsFormat(data, group_ratio);
     } else {
       showError(message);
@@ -307,14 +349,14 @@ const ModelPricing = () => {
             type="success"
             fullMode={false}
             closeIcon="null"
-            description={`您的分组为：${userState.user.group}，分组倍率为：${groupRatio}`}
+            description={`您的默认分组为：${userState.user.group}，分组倍率为：${groupRatio[userState.user.group]}`}
           />
         ) : (
           <Banner
             type='warning'
             fullMode={false}
             closeIcon="null"
-            description={`您还未登陆，显示的价格为默认分组倍率: ${groupRatio}`}
+            description={`您还未登陆，显示的价格为默认分组倍率: ${groupRatio['default']}`}
           />
         )}
         <br/>

web/src/pages/Token/EditToken.js

@@ -273,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>
                 启用模型限制（非必要，不建议启用）
@@ -425,27 +425,27 @@ 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 style={{ marginTop: 10 }}>
+            <Typography.Text>令牌分组，默认为用户的分组</Typography.Text>
           </div>
-          <Select
-              style={{marginTop: 8}}
-              placeholder={'令牌分组，不选为默认分组'}
+          {groups.length > 0 ?
+            <Select
+              style={{ marginTop: 8 }}
+              placeholder={'令牌分组，默认为用户的分组'}
               name='gruop'
               required
               selection
@@ -455,7 +455,14 @@ const EditToken = (props) => {
               value={inputs.group}
               autoComplete='new-password'
               optionList={groups}
-          />
+            />:
+            <Select
+              style={{ marginTop: 8 }}
+              placeholder={'管理员未设置用户可选分组'}
+              name='gruop'
+              disabled={true}
+            />
+          }
         </Spin>
       </SideSheet>
     </>