feat: aiproxy modelinfo and dashboard and model rpm limit

service/aiproxy/common/config/config.go

@@ -1,144 +1,129 @@
 package config
 
 import (
+	"math"
 	"os"
 	"slices"
-	"strconv"
-	"sync"
 	"sync/atomic"
-	"time"
 
 	"github.com/labring/sealos/service/aiproxy/common/env"
 )
 
 var (
-	OptionMap        map[string]string
-	OptionMapRWMutex sync.RWMutex
+	DebugEnabled    = env.Bool("DEBUG", false)
+	DebugSQLEnabled = env.Bool("DEBUG_SQL", false)
 )
 
 var (
-	DebugEnabled, _    = strconv.ParseBool(os.Getenv("DEBUG"))
-	DebugSQLEnabled, _ = strconv.ParseBool(os.Getenv("DEBUG_SQL"))
+	DisableAutoMigrateDB = env.Bool("DISABLE_AUTO_MIGRATE_DB", false)
+	OnlyOneLogFile       = env.Bool("ONLY_ONE_LOG_FILE", false)
+	AdminKey             = os.Getenv("ADMIN_KEY")
 )
 
 var (
-	// 当测试或请求的时候发生错误是否自动禁用渠道
-	automaticDisableChannelEnabled atomic.Bool
-	// 当测试成功是否自动启用渠道
-	automaticEnableChannelWhenTestSucceedEnabled atomic.Bool
-	// 是否近似计算token
-	approximateTokenEnabled atomic.Bool
-	// 重试次数
-	retryTimes atomic.Int64
-	// 暂停服务
-	disableServe atomic.Bool
-	// log detail 存储时间(小时)
+	disableServe          atomic.Bool
 	logDetailStorageHours int64 = 3 * 24
+	internalToken         atomic.Value
 )
 
-func GetLogDetailStorageHours() int64 {
-	return atomic.LoadInt64(&logDetailStorageHours)
-}
-
-func SetLogDetailStorageHours(hours int64) {
-	atomic.StoreInt64(&logDetailStorageHours, hours)
-}
-
-func GetDisableServe() bool {
-	return disableServe.Load()
-}
-
-func SetDisableServe(disabled bool) {
-	disableServe.Store(disabled)
-}
+var (
+	retryTimes              atomic.Int64
+	enableModelErrorAutoBan atomic.Bool
+	modelErrorAutoBanRate   = math.Float64bits(0.5)
+	timeoutWithModelType    atomic.Value
+	disableModelConfig      atomic.Bool
+)
 
-func GetAutomaticDisableChannelEnabled() bool {
-	return automaticDisableChannelEnabled.Load()
-}
+var (
+	defaultChannelModels       atomic.Value
+	defaultChannelModelMapping atomic.Value
+	groupMaxTokenNum           atomic.Int64
+	groupConsumeLevelRatio     atomic.Value
+)
 
-func SetAutomaticDisableChannelEnabled(enabled bool) {
-	automaticDisableChannelEnabled.Store(enabled)
-}
+var geminiSafetySetting atomic.Value
 
-func GetAutomaticEnableChannelWhenTestSucceedEnabled() bool {
-	return automaticEnableChannelWhenTestSucceedEnabled.Load()
-}
+var billingEnabled atomic.Bool
 
-func SetAutomaticEnableChannelWhenTestSucceedEnabled(enabled bool) {
-	automaticEnableChannelWhenTestSucceedEnabled.Store(enabled)
+func init() {
+	timeoutWithModelType.Store(make(map[int]int64))
+	defaultChannelModels.Store(make(map[int][]string))
+	defaultChannelModelMapping.Store(make(map[int]map[string]string))
+	groupConsumeLevelRatio.Store(make(map[float64]float64))
+	geminiSafetySetting.Store("BLOCK_NONE")
+	billingEnabled.Store(true)
+	internalToken.Store(os.Getenv("INTERNAL_TOKEN"))
 }
 
-func GetApproximateTokenEnabled() bool {
-	return approximateTokenEnabled.Load()
+func GetDisableModelConfig() bool {
+	return disableModelConfig.Load()
 }
 
-func SetApproximateTokenEnabled(enabled bool) {
-	approximateTokenEnabled.Store(enabled)
+func SetDisableModelConfig(disabled bool) {
+	disabled = env.Bool("DISABLE_MODEL_CONFIG", disabled)
+	disableModelConfig.Store(disabled)
 }
 
 func GetRetryTimes() int64 {
 	return retryTimes.Load()
 }
 
 func SetRetryTimes(times int64) {
+	times = env.Int64("RETRY_TIMES", times)
 	retryTimes.Store(times)
 }
 
-var DisableAutoMigrateDB = os.Getenv("DISABLE_AUTO_MIGRATE_DB") == "true"
-
-var RelayTimeout = env.Int("RELAY_TIMEOUT", 0) // unit is second
-
-var RateLimitKeyExpirationDuration = 20 * time.Minute
+func GetEnableModelErrorAutoBan() bool {
+	return enableModelErrorAutoBan.Load()
+}
 
-var OnlyOneLogFile = env.Bool("ONLY_ONE_LOG_FILE", false)
+func SetEnableModelErrorAutoBan(enabled bool) {
+	enabled = env.Bool("ENABLE_MODEL_ERROR_AUTO_BAN", enabled)
+	enableModelErrorAutoBan.Store(enabled)
+}
 
-var (
-	// 代理地址
-	RelayProxy = env.String("RELAY_PROXY", "")
-	// 用户内容请求代理地址
-	UserContentRequestProxy = env.String("USER_CONTENT_REQUEST_PROXY", "")
-	// 用户内容请求超时时间，单位为秒
-	UserContentRequestTimeout = env.Int("USER_CONTENT_REQUEST_TIMEOUT", 30)
-)
+func GetModelErrorAutoBanRate() float64 {
+	return math.Float64frombits(atomic.LoadUint64(&modelErrorAutoBanRate))
+}
 
-var AdminKey = env.String("ADMIN_KEY", "")
+func SetModelErrorAutoBanRate(rate float64) {
+	rate = env.Float64("MODEL_ERROR_AUTO_BAN_RATE", rate)
+	atomic.StoreUint64(&modelErrorAutoBanRate, math.Float64bits(rate))
+}
 
-var (
-	globalAPIRateLimitNum      atomic.Int64
-	defaultChannelModels       atomic.Value
-	defaultChannelModelMapping atomic.Value
-	defaultGroupQPM            atomic.Int64
-	groupMaxTokenNum           atomic.Int32
-)
+func GetTimeoutWithModelType() map[int]int64 {
+	return timeoutWithModelType.Load().(map[int]int64)
+}
 
-func init() {
-	defaultChannelModels.Store(make(map[int][]string))
-	defaultChannelModelMapping.Store(make(map[int]map[string]string))
+func SetTimeoutWithModelType(timeout map[int]int64) {
+	timeout = env.JSON("TIMEOUT_WITH_MODEL_TYPE", timeout)
+	timeoutWithModelType.Store(timeout)
 }
 
-// 全局qpm，不是根据ip限制，而是所有请求共享一个qpm
-func GetGlobalAPIRateLimitNum() int64 {
-	return globalAPIRateLimitNum.Load()
+func GetLogDetailStorageHours() int64 {
+	return atomic.LoadInt64(&logDetailStorageHours)
 }
 
-func SetGlobalAPIRateLimitNum(num int64) {
-	globalAPIRateLimitNum.Store(num)
+func SetLogDetailStorageHours(hours int64) {
+	hours = env.Int64("LOG_DETAIL_STORAGE_HOURS", hours)
+	atomic.StoreInt64(&logDetailStorageHours, hours)
 }
 
-// group默认qpm，如果group没有设置qpm，则使用该qpm
-func GetDefaultGroupQPM() int64 {
-	return defaultGroupQPM.Load()
+func GetDisableServe() bool {
+	return disableServe.Load()
 }
 
-func SetDefaultGroupQPM(qpm int64) {
-	defaultGroupQPM.Store(qpm)
+func SetDisableServe(disabled bool) {
+	disabled = env.Bool("DISABLE_SERVE", disabled)
+	disableServe.Store(disabled)
 }
 
 func GetDefaultChannelModels() map[int][]string {
 	return defaultChannelModels.Load().(map[int][]string)
 }
 
 func SetDefaultChannelModels(models map[int][]string) {
+	models = env.JSON("DEFAULT_CHANNEL_MODELS", models)
 	for key, ms := range models {
 		slices.Sort(ms)
 		models[key] = slices.Compact(ms)
@@ -151,54 +136,52 @@ func GetDefaultChannelModelMapping() map[int]map[string]string {
 }
 
 func SetDefaultChannelModelMapping(mapping map[int]map[string]string) {
+	mapping = env.JSON("DEFAULT_CHANNEL_MODEL_MAPPING", mapping)
 	defaultChannelModelMapping.Store(mapping)
 }
 
-// 那个group最多可创建的token数量，0表示不限制
-func GetGroupMaxTokenNum() int32 {
-	return groupMaxTokenNum.Load()
+func GetGroupConsumeLevelRatio() map[float64]float64 {
+	return groupConsumeLevelRatio.Load().(map[float64]float64)
 }
 
-func SetGroupMaxTokenNum(num int32) {
-	groupMaxTokenNum.Store(num)
+func SetGroupConsumeLevelRatio(ratio map[float64]float64) {
+	ratio = env.JSON("GROUP_CONSUME_LEVEL_RATIO", ratio)
+	groupConsumeLevelRatio.Store(ratio)
 }
 
-var (
-	geminiSafetySetting atomic.Value
-	geminiVersion       atomic.Value
-)
+// GetGroupMaxTokenNum returns max number of tokens per group, 0 means unlimited
+func GetGroupMaxTokenNum() int64 {
+	return groupMaxTokenNum.Load()
+}
 
-func init() {
-	geminiSafetySetting.Store("BLOCK_NONE")
-	geminiVersion.Store("v1beta")
+func SetGroupMaxTokenNum(num int64) {
+	num = env.Int64("GROUP_MAX_TOKEN_NUM", num)
+	groupMaxTokenNum.Store(num)
 }
 
 func GetGeminiSafetySetting() string {
 	return geminiSafetySetting.Load().(string)
 }
 
 func SetGeminiSafetySetting(setting string) {
+	setting = env.String("GEMINI_SAFETY_SETTING", setting)
 	geminiSafetySetting.Store(setting)
 }
 
-func GetGeminiVersion() string {
-	return geminiVersion.Load().(string)
-}
-
-func SetGeminiVersion(version string) {
-	geminiVersion.Store(version)
-}
-
-var billingEnabled atomic.Bool
-
-func init() {
-	billingEnabled.Store(true)
-}
-
 func GetBillingEnabled() bool {
 	return billingEnabled.Load()
 }
 
 func SetBillingEnabled(enabled bool) {
+	enabled = env.Bool("BILLING_ENABLED", enabled)
 	billingEnabled.Store(enabled)
 }
+
+func GetInternalToken() string {
+	return internalToken.Load().(string)
+}
+
+func SetInternalToken(token string) {
+	token = env.String("INTERNAL_TOKEN", token)
+	internalToken.Store(token)
+}