fix: '/vi/models' #474

controller/model.go

@@ -146,9 +146,35 @@ func ListModels(c *gin.Context) {
 		})
 		})
 		return
 		return
 	}
 	}
-	models := model.GetGroupModels(user.Group)
 	userOpenAiModels := make([]dto.OpenAIModels, 0)
 	userOpenAiModels := make([]dto.OpenAIModels, 0)
 	permission := getPermission()
 	permission := getPermission()
+
+	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{}
+		}
+		for allowModel, _ := range tokenModelLimit {
+			if _, ok := openAIModelsMap[allowModel]; ok {
+				userOpenAiModels = append(userOpenAiModels, openAIModelsMap[allowModel])
+			} else {
+				userOpenAiModels = append(userOpenAiModels, dto.OpenAIModels{
+					Id:         allowModel,
+					Object:     "model",
+					Created:    1626777600,
+					OwnedBy:    "custom",
+					Permission: permission,
+					Root:       allowModel,
+					Parent:     nil,
+				})
+			}
+		}
+	} else {
+		models := model.GetGroupModels(user.Group)
 		for _, s := range models {
 		for _, s := range models {
 			if _, ok := openAIModelsMap[s]; ok {
 			if _, ok := openAIModelsMap[s]; ok {
 				userOpenAiModels = append(userOpenAiModels, openAIModelsMap[s])
 				userOpenAiModels = append(userOpenAiModels, openAIModelsMap[s])
@@ -164,6 +190,7 @@ func ListModels(c *gin.Context) {
 				})
 				})
 			}
 			}
 		}
 		}
+	}
 	c.JSON(200, gin.H{
 	c.JSON(200, gin.H{
 		"success": true,
 		"success": true,
 		"data":    userOpenAiModels,
 		"data":    userOpenAiModels,