tbphp 10 месяцев назад
Родитель
Сommit
88ed83f419
2 измененных файлов с 25 добавлено и 0 удалено
  1. 9 0
      controller/option.go
  2. 16 0
      setting/rate_limit.go

+ 9 - 0
controller/option.go

@@ -110,6 +110,15 @@ func UpdateOption(c *gin.Context) {
 			})
 			return
 		}
+	case "ModelRequestRateLimitGroup":
+		err = setting.CheckModelRequestRateLimitGroup(option.Value)
+		if err != nil {
+			c.JSON(http.StatusOK, gin.H{
+				"success": false,
+				"message": err.Error(),
+			})
+			return
+		}
 
 	}
 	err = model.UpdateOption(option.Key, option.Value)

+ 16 - 0
setting/rate_limit.go

@@ -2,6 +2,7 @@ package setting
 
 import (
 	"encoding/json"
+	"fmt"
 	"one-api/common"
 	"sync"
 )
@@ -46,3 +47,18 @@ func GetGroupRateLimit(group string) (totalCount, successCount int, found bool)
 	}
 	return limits[0], limits[1], true
 }
+
+func CheckModelRequestRateLimitGroup(jsonStr string) error {
+	checkModelRequestRateLimitGroup := make(map[string][2]int)
+	err := json.Unmarshal([]byte(jsonStr), &checkModelRequestRateLimitGroup)
+	if err != nil {
+		return err
+	}
+	for group, limits := range checkModelRequestRateLimitGroup {
+		if limits[0] < 0 || limits[1] < 0 {
+			return fmt.Errorf("group %s has negative rate limit values: [%d, %d]", group, limits[0], limits[1])
+		}
+	}
+
+	return nil
+}