featL add token-side model selection

controller/relay-text.go

@@ -69,6 +69,27 @@ func relayTextHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode {
 			isModelMapped = true
 		}
 	}
+
+	// Get token info
+	tokenInfo, err := model.GetTokenById(tokenId)
+
+	if err != nil {
+		return errorWrapper(err, "get_token_info_failed", http.StatusInternalServerError)
+	}
+
+	hasModelAvailable := func() bool {
+		for _, token := range strings.Split(tokenInfo.Models, ",") {
+			if token == textRequest.Model {
+				return true
+			}
+		}
+		return false
+	}()
+
+	if !hasModelAvailable {
+		return errorWrapper(errors.New("model not available for use"), "model_not_available_for_use", http.StatusBadRequest)
+	}
+
 	baseURL := common.ChannelBaseURLs[channelType]
 	requestURL := c.Request.URL.String()
 	if c.GetString("base_url") != "" {

controller/token.go

@@ -1,11 +1,12 @@
 package controller
 
 import (
-	"github.com/gin-gonic/gin"
 	"net/http"
 	"one-api/common"
 	"one-api/model"
 	"strconv"
+
+	"github.com/gin-gonic/gin"
 )
 
 func GetAllTokens(c *gin.Context) {
@@ -203,6 +204,7 @@ func UpdateToken(c *gin.Context) {
 		cleanToken.ExpiredTime = token.ExpiredTime
 		cleanToken.RemainQuota = token.RemainQuota
 		cleanToken.UnlimitedQuota = token.UnlimitedQuota
+		cleanToken.Models = token.Models
 	}
 	err = cleanToken.Update()
 	if err != nil {

i18n/en.json

@@ -523,5 +523,6 @@
   "该 Discord 账户已被绑定": "The Discord account has been bound",
   "管理员未开启通过 Discord 登录以及注册": "The administrator has not enabled login and registration via Discord",
   "无法启用 Discord OAuth，请先填入 Discord Client ID 以及 Discord Client Secret！": "Unable to enable Discord OAuth, please fill in the Discord Client ID and Discord Client Secret first!",
-  "兑换失败，": "Redemption failed, "
+  "兑换失败，": "Redemption failed, ",
+  "请选择此密钥支持的模型": "Please select the models supported by this key"
 }

model/token.go

@@ -3,8 +3,9 @@ package model
 import (
 	"errors"
 	"fmt"
-	"gorm.io/gorm"
 	"one-api/common"
+
+	"gorm.io/gorm"
 )
 
 type Token struct {
@@ -19,6 +20,7 @@ type Token struct {
 	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
+	Models         string `json:"models"`
 }
 
 func GetAllUserTokens(userId int, startIdx int, num int) ([]*Token, error) {
@@ -99,7 +101,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", "models").Updates(token).Error
 	return err
 }
 

web/src/pages/Token/EditToken.js

@@ -13,8 +13,12 @@ const EditToken = () => {
     name: '',
     remain_quota: isEdit ? 0 : 500000,
     expired_time: -1,
-    unlimited_quota: false
+    unlimited_quota: false,
+    models: [],
   };
+  const [modelOptions, setModelOptions] = useState([]);
+  const [basicModels, setBasicModels] = useState([]);
+  const [fullModels, setFullModels] = useState([]);
   const [inputs, setInputs] = useState(originInputs);
   const { name, remain_quota, expired_time, unlimited_quota } = inputs;
 
@@ -41,6 +45,21 @@ const EditToken = () => {
     setInputs({ ...inputs, unlimited_quota: !unlimited_quota });
   };
 
+  const fetchModels = async () => {
+    try {
+      let res = await API.get(`/api/channel/models`);
+      setModelOptions(res.data.data.map((model) => ({
+        key: model.id,
+        text: model.id,
+        value: model.id
+      })));
+      setFullModels(res.data.data.map((model) => model.id));
+      setBasicModels(res.data.data.filter((model) => !model.id.startsWith('gpt-4')).map((model) => model.id));
+    } catch (error) {
+      showError(error.message);
+    }
+  };
+
   const loadToken = async () => {
     let res = await API.get(`/api/token/${tokenId}`);
     const { success, message, data } = res.data;
@@ -72,6 +91,11 @@ const EditToken = () => {
       }
       localInputs.expired_time = Math.ceil(time / 1000);
     }
+    if (inputs.models.length === 0) {
+      showError('请至少选择一个模型！');
+      return;
+    }
+    localInputs.models = localInputs.models.join(',');
     let res;
     if (isEdit) {
       res = await API.put(`/api/token/`, { ...localInputs, id: parseInt(tokenId) });
@@ -91,6 +115,10 @@ const EditToken = () => {
     }
   };
 
+  useEffect(() => {
+    fetchModels().then();
+  }, []);
+
   return (
     <>
       <Segment loading={loading}>
@@ -151,6 +179,32 @@ const EditToken = () => {
           <Button type={'button'} onClick={() => {
             setUnlimitedQuota();
           }}>{unlimited_quota ? '取消无限额度' : '设置为无限额度'}</Button>
+          <Form.Field style={{ marginTop: '12px' }}>
+            <Form.Dropdown
+              label='模型'
+              placeholder={'请选择此密钥支持的模型'}
+              name='models'
+              required
+              fluid
+              multiple
+              selection
+              onChange={handleInputChange}
+              value={inputs.models}
+              autoComplete='new-password'
+              options={modelOptions}
+            />
+          </Form.Field>
+          <div style={{ lineHeight: '40px', marginBottom: '12px' }}>
+            <Button type={'button'} onClick={() => {
+              handleInputChange(null, { name: 'models', value: basicModels });
+            }}>填入基础模型</Button>
+            <Button type={'button'} onClick={() => {
+              handleInputChange(null, { name: 'models', value: fullModels });
+            }}>填入所有模型</Button>
+            <Button type={'button'} onClick={() => {
+              handleInputChange(null, { name: 'models', value: [] });
+            }}>清除所有模型</Button>
+          </div>
           <Button positive onClick={submit}>提交</Button>
         </Form>
       </Segment>