Merge pull request #1267 from t0ng7u/feature/upstream-ratio-sync

🔄 feat(ratio-sync): introduce upstream ratio synchronisation feature #1220

setting/ratio_setting/cache_ratio.go

@@ -85,7 +85,11 @@ func UpdateCacheRatioByJSONString(jsonStr string) error {
 	cacheRatioMapMutex.Lock()
 	defer cacheRatioMapMutex.Unlock()
 	cacheRatioMap = make(map[string]float64)
-	return json.Unmarshal([]byte(jsonStr), &cacheRatioMap)
+	err := json.Unmarshal([]byte(jsonStr), &cacheRatioMap)
+	if err == nil {
+		InvalidateExposedDataCache()
+	}
+	return err
 }
 
 // GetCacheRatio returns the cache ratio for a model
@@ -106,3 +110,13 @@ func GetCreateCacheRatio(name string) (float64, bool) {
 	}
 	return ratio, true
 }
+
+func GetCacheRatioCopy() map[string]float64 {
+	cacheRatioMapMutex.RLock()
+	defer cacheRatioMapMutex.RUnlock()
+	copyMap := make(map[string]float64, len(cacheRatioMap))
+	for k, v := range cacheRatioMap {
+		copyMap[k] = v
+	}
+	return copyMap
+}

setting/ratio_setting/expose_ratio.go

@@ -0,0 +1,17 @@
+package ratio_setting
+
+import "sync/atomic"
+
+var exposeRatioEnabled atomic.Bool
+
+func init() {
+    exposeRatioEnabled.Store(false)
+}
+
+func SetExposeRatioEnabled(enabled bool) {
+    exposeRatioEnabled.Store(enabled)
+}
+
+func IsExposeRatioEnabled() bool {
+    return exposeRatioEnabled.Load()
+} 

setting/ratio_setting/exposed_cache.go

@@ -0,0 +1,55 @@
+package ratio_setting
+
+import (
+    "sync"
+    "sync/atomic"
+    "time"
+
+    "github.com/gin-gonic/gin"
+)
+
+const exposedDataTTL = 30 * time.Second
+
+type exposedCache struct {
+    data      gin.H
+    expiresAt time.Time
+}
+
+var (
+    exposedData atomic.Value
+    rebuildMu   sync.Mutex
+)
+
+func InvalidateExposedDataCache() {
+    exposedData.Store((*exposedCache)(nil))
+}
+
+func cloneGinH(src gin.H) gin.H {
+    dst := make(gin.H, len(src))
+    for k, v := range src {
+        dst[k] = v
+    }
+    return dst
+}
+
+func GetExposedData() gin.H {
+    if c, ok := exposedData.Load().(*exposedCache); ok && c != nil && time.Now().Before(c.expiresAt) {
+        return cloneGinH(c.data)
+    }
+    rebuildMu.Lock()
+    defer rebuildMu.Unlock()
+    if c, ok := exposedData.Load().(*exposedCache); ok && c != nil && time.Now().Before(c.expiresAt) {
+        return cloneGinH(c.data)
+    }
+    newData := gin.H{
+        "model_ratio":      GetModelRatioCopy(),
+        "completion_ratio": GetCompletionRatioCopy(),
+        "cache_ratio":      GetCacheRatioCopy(),
+        "model_price":      GetModelPriceCopy(),
+    }
+    exposedData.Store(&exposedCache{
+        data:      newData,
+        expiresAt: time.Now().Add(exposedDataTTL),
+    })
+    return cloneGinH(newData)
+} 

setting/ratio_setting/model_ratio.go

@@ -317,7 +317,11 @@ func UpdateModelPriceByJSONString(jsonStr string) error {
 	modelPriceMapMutex.Lock()
 	defer modelPriceMapMutex.Unlock()
 	modelPriceMap = make(map[string]float64)
-	return json.Unmarshal([]byte(jsonStr), &modelPriceMap)
+	err := json.Unmarshal([]byte(jsonStr), &modelPriceMap)
+	if err == nil {
+		InvalidateExposedDataCache()
+	}
+	return err
 }
 
 // GetModelPrice 返回模型的价格，如果模型不存在则返回-1，false
@@ -345,7 +349,11 @@ func UpdateModelRatioByJSONString(jsonStr string) error {
 	modelRatioMapMutex.Lock()
 	defer modelRatioMapMutex.Unlock()
 	modelRatioMap = make(map[string]float64)
-	return json.Unmarshal([]byte(jsonStr), &modelRatioMap)
+	err := json.Unmarshal([]byte(jsonStr), &modelRatioMap)
+	if err == nil {
+		InvalidateExposedDataCache()
+	}
+	return err
 }
 
 // 处理带有思考预算的模型名称，方便统一定价
@@ -405,7 +413,11 @@ func UpdateCompletionRatioByJSONString(jsonStr string) error {
 	CompletionRatioMutex.Lock()
 	defer CompletionRatioMutex.Unlock()
 	CompletionRatio = make(map[string]float64)
-	return json.Unmarshal([]byte(jsonStr), &CompletionRatio)
+	err := json.Unmarshal([]byte(jsonStr), &CompletionRatio)
+	if err == nil {
+		InvalidateExposedDataCache()
+	}
+	return err
 }
 
 func GetCompletionRatio(name string) float64 {
@@ -609,3 +621,33 @@ func GetImageRatio(name string) (float64, bool) {
 	}
 	return ratio, true
 }
+
+func GetModelRatioCopy() map[string]float64 {
+	modelRatioMapMutex.RLock()
+	defer modelRatioMapMutex.RUnlock()
+	copyMap := make(map[string]float64, len(modelRatioMap))
+	for k, v := range modelRatioMap {
+		copyMap[k] = v
+	}
+	return copyMap
+}
+
+func GetModelPriceCopy() map[string]float64 {
+	modelPriceMapMutex.RLock()
+	defer modelPriceMapMutex.RUnlock()
+	copyMap := make(map[string]float64, len(modelPriceMap))
+	for k, v := range modelPriceMap {
+		copyMap[k] = v
+	}
+	return copyMap
+}
+
+func GetCompletionRatioCopy() map[string]float64 {
+	CompletionRatioMutex.RLock()
+	defer CompletionRatioMutex.RUnlock()
+	copyMap := make(map[string]float64, len(CompletionRatio))
+	for k, v := range CompletionRatio {
+		copyMap[k] = v
+	}
+	return copyMap
+}