feat: #15 支持不修改数据库和强制使用内置模型

Author: DullJZ

README-zh.md

@@ -87,11 +87,13 @@ docker-compose up -d
   "exclude_model": ["advanced-voice", "minimax_s2v-01", "minimax_video-01", "minimax_video-01-live2d"],
   "models": ["gpt-3.5-turbo", "gpt-4o"],
   "force_models": false,
+  "force_inside_models": false,
   "time_period": "1h",
   "max_concurrent": 5,
   "rps": 5,
   "db_type": "YOUR_DB_TYPE",
   "db_dsn": "YOUR_DB_DSN",
+  "do_not_modify_db": false,
   "base_url": "http://localhost:3000",
   "system_token": "YOUR_SYSTEM_TOKEN",
   "uptime-kuma": {
@@ -133,11 +135,13 @@ docker-compose up -d
 - exclude_model: 排除不予监控的模型ID  
 - models: 模型列表，仅当获取不到渠道的模型(/v1/models)时使用
 - force_models: 如果为true，将强制只测试上述模型，不再获取渠道的模型，默认为false
+- force_inside_models: 如果为true，将强制只测试OneAPI设置的模型，不再获取模型列表，默认为false。如果force_models为true，此项无效 
 - time_period: 模型可用性测试的时间间隔，建议不小于30分钟，接收的时间格式为s、m、h
 - max_concurrency: 在一个渠道内测试的最大并发数，默认为5
 - rps: 在一个渠道内测试的每秒请求数，默认为5
 - db_type: 数据库类型，包括mysql、sqlite、postgres、sqlserver
 - db_dsn: 数据库DSN字符串，不同数据库类型的DSN格式不同，示例如下
+- do_not_modify_db: 如果为true，将不会修改数据库中的可用模型，默认为false
 - base_url: OneAPI/NewAPI/OneHub的基础URL，如果使用host模式，可以直接使用http://localhost:3000，目前只有OneHub需要填写
 - system_token: 系统Token，目前只有OneHub需要填写
 - uptime-kuma: Uptime Kuma的配置，status为`enabled`或`disabled`，model_url和channel_url为模型和渠道的可用性Push URL

README.md

@@ -86,11 +86,13 @@ The configuration file is `config.json` located in the same directory, with the
   "exclude_model": ["advanced-voice", "minimax_s2v-01", "minimax_video-01", "minimax_video-01-live2d"],
   "models": ["gpt-3.5-turbo", "gpt-4o"],
   "force_models": false,
+  "force_inside_models": false,
   "time_period": "1h",
   "max_concurrent": 5,
   "rps": 5,
   "db_type": "YOUR_DB_TYPE",
   "db_dsn": "YOUR_DB_DSN",
+  "do_not_modify_db": false,
   "base_url": "http://localhost:3000",
   "system_token": "YOUR_SYSTEM_TOKEN",
   "uptime-kuma": {
@@ -132,11 +134,13 @@ Configuration explanation:
 - exclude_model: IDs of models to exclude from monitoring
 - models: List of models, used only when unable to retrieve models from the channel (/v1/models)
 - force_models: If true, only the above models will be tested, and channel models will not be fetched. Default is false
+- force_inside_models: If true, only the models set in OneAPI will be tested, and the model list will not be fetched. Default is false. If force_models is true, this option is invalid.
 - time_period: Interval for testing model availability, recommended not less than 30 minutes, accepts time formats s, m, h
 - max_concurrency: Maximum number of concurrent tests within a channel, default is 5
 - rps: Requests per second within a channel, default is 5
 - db_type: Database type, including mysql, sqlite, postgres, sqlserver
 - db_dsn: Database DSN string, the format varies by database type. Examples below
+- do_not_modify_db: If true, the available models in the database will not be modified. Default is false
 - base_url: The base URL for OneAPI/NewAPI/OneHub. If using host mode, you can directly use http://localhost:3000. Currently, only OneHub requires this field.
 - system_token: System token, currently only required for OneHub.
 - uptime-kuma: Configuration for Uptime Kuma. The status can be `enabled` or `disabled`. The model_url and channel_url are the availability Push URLs for models and channels.

VERSION

@@ -1 +1 @@
-v0.4.5
+v0.4.6

config.go

@@ -7,19 +7,21 @@ import (
 )
 
 type Config struct {
-	OneAPIType     string   `json:"oneapi_type"`
-	ExcludeChannel []int    `json:"exclude_channel"`
-	ExcludeModel   []string `json:"exclude_model"`
-	Models         []string `json:"models"`
-	ForceModels    bool     `json:"force_models"`
-	TimePeriod     string   `json:"time_period"`
-	MaxConcurrent  int      `json:"max_concurrent"`
-	RPS            int      `json:"rps"`
-	DbType         string   `json:"db_type"`
-	DbDsn          string   `json:"db_dsn"`
-	BaseURL        string   `json:"base_url"`
-	SystemToken    string   `json:"system_token"`
-	UptimeKuma     struct {
+	OneAPIType        string   `json:"oneapi_type"`
+	ExcludeChannel    []int    `json:"exclude_channel"`
+	ExcludeModel      []string `json:"exclude_model"`
+	Models            []string `json:"models"`
+	ForceModels       bool     `json:"force_models"`
+	ForceInsideModels bool     `json:"force_inside_models"`
+	TimePeriod        string   `json:"time_period"`
+	MaxConcurrent     int      `json:"max_concurrent"`
+	RPS               int      `json:"rps"`
+	DbType            string   `json:"db_type"`
+	DbDsn             string   `json:"db_dsn"`
+	DoNotModifyDb     bool     `json:"do_not_modify_db"`
+	BaseURL           string   `json:"base_url"`
+	SystemToken       string   `json:"system_token"`
+	UptimeKuma        struct {
 		Status     string            `json:"status"`
 		ModelURL   map[string]string `json:"model_url"`
 		ChannelURL map[string]string `json:"channel_url"`

config_example.json

@@ -4,11 +4,13 @@
     "exclude_model": ["advanced-voice", "minimax_s2v-01", "minimax_video-01", "minimax_video-01-live2d"],
     "models": ["gpt-3.5-turbo", "gpt-4o"],
     "force_models": false,
+    "force_inside_models": false,
     "time_period": "1h",
     "max_concurrent": 5,
     "rps": 5,
     "db_type": "YOUR_DB_TYPE",
     "db_dsn": "YOUR_DB_DSN",
+    "do_not_modify_db": false,
     "base_url": "http://localhost:3000",
     "system_token": "YOUR_SYSTEM_TOKEN",
     "uptime-kuma": {

main.go

@@ -101,45 +101,56 @@ func testModels(channel Channel, wg *sync.WaitGroup, mu *sync.Mutex) {
 		log.Println("强制使用自定义模型列表")
 		modelList = config.Models
 	} else {
-		// 从/v1/models接口获取模型列表
-		req, err := http.NewRequest("GET", channel.BaseURL+"/v1/models", nil)
-		if err != nil {
-			log.Printf("创建请求失败：%v\n", err)
-			return
-		}
-		req.Header.Set("Authorization", "Bearer "+channel.Key)
-
-		client := &http.Client{}
-		resp, err := client.Do(req)
-		if err != nil {
-			log.Println("获取模型列表失败：", err, "尝试自定义模型列表")
-			modelList = config.Models
+		if config.ForceInsideModels {
+			log.Println("强制使用内置模型列表")
+			// 从数据库获取模型列表
+			var models string
+			if err := db.Raw("SELECT models FROM channels WHERE id = ?", channel.ID).Scan(&models).Error; err != nil {
+				log.Printf("获取渠道 %s(ID:%d) 的模型列表失败：%v\n", channel.Name, channel.ID, err)
+				return
+			}
+			modelList = strings.Split(models, ",")
 		} else {
-			defer resp.Body.Close()
-			body, _ := ioutil.ReadAll(resp.Body)
-			if resp.StatusCode != http.StatusOK {
-				log.Printf("获取模型列表失败，状态码：%d，响应：%s\n", resp.StatusCode, string(body))
+			// 从/v1/models接口获取模型列表
+			req, err := http.NewRequest("GET", channel.BaseURL+"/v1/models", nil)
+			if err != nil {
+				log.Printf("创建请求失败：%v\n", err)
 				return
 			}
+			req.Header.Set("Authorization", "Bearer "+channel.Key)
 
-			// 解析响应JSON
-			var response struct {
-				Data []struct {
-					ID string `json:"id"`
-				} `json:"data"`
-			}
+			client := &http.Client{}
+			resp, err := client.Do(req)
+			if err != nil {
+				log.Println("获取模型列表失败：", err, "尝试自定义模型列表")
+				modelList = config.Models
+			} else {
+				defer resp.Body.Close()
+				body, _ := ioutil.ReadAll(resp.Body)
+				if resp.StatusCode != http.StatusOK {
+					log.Printf("获取模型列表失败，状态码：%d，响应：%s\n", resp.StatusCode, string(body))
+					return
+				}
 
-			if err := json.Unmarshal(body, &response); err != nil {
-				log.Printf("解析模型列表失败：%v\n", err)
-				return
-			}
-			// 提取模型ID列表
-			for _, model := range response.Data {
-				if containsString(config.ExcludeModel, model.ID) {
-					log.Printf("模型 %s 在排除列表中，跳过\n", model.ID)
-					continue
+				// 解析响应JSON
+				var response struct {
+					Data []struct {
+						ID string `json:"id"`
+					} `json:"data"`
+				}
+
+				if err := json.Unmarshal(body, &response); err != nil {
+					log.Printf("解析模型列表失败：%v\n", err)
+					return
+				}
+				// 提取模型ID列表
+				for _, model := range response.Data {
+					if containsString(config.ExcludeModel, model.ID) {
+						log.Printf("模型 %s 在排除列表中，跳过\n", model.ID)
+						continue
+					}
+					modelList = append(modelList, model.ID)
 				}
-				modelList = append(modelList, model.ID)
 			}
 		}
 	}
@@ -220,6 +231,10 @@ func testModels(channel Channel, wg *sync.WaitGroup, mu *sync.Mutex) {
 	modelWg.Wait()
 
 	// 更新模型
+	if config.DoNotModifyDb {
+		log.Println("跳过数据库更新")
+		return
+	}
 	mu.Lock()
 	err := updateModels(channel.ID, availableModels, channel.ModelMapping)
 	mu.Unlock()