From 9d9c461c48876f0bb503c6446f8b31b4c4a2d9dc Mon Sep 17 00:00:00 2001
From: "1808837298@qq.com" <1808837298@qq.com>
Date: Tue, 18 Feb 2025 15:01:43 +0800
Subject: [PATCH] refactor: Improve CompletionRatio handling with thread-safe
 access and initialization

---
 common/model-ratio.go | 29 +++++++++++++++--------------
 1 file changed, 15 insertions(+), 14 deletions(-)

diff --git a/common/model-ratio.go b/common/model-ratio.go
index ffeda83d..4b64c79f 100644
--- a/common/model-ratio.go
+++ b/common/model-ratio.go
@@ -233,7 +233,11 @@ var (
 	modelRatioMapMutex                    = sync.RWMutex{}
 )
 
-var CompletionRatio map[string]float64 = nil
+var (
+	CompletionRatio      map[string]float64 = nil
+	CompletionRatioMutex                    = sync.RWMutex{}
+)
+
 var defaultCompletionRatio = map[string]float64{
 	"gpt-4-gizmo-*":  2,
 	"gpt-4o-gizmo-*": 3,
@@ -334,10 +338,17 @@ func GetDefaultModelRatioMap() map[string]float64 {
 	return defaultModelRatio
 }
 
-func CompletionRatio2JSONString() string {
+func GetCompletionRatioMap() map[string]float64 {
+	CompletionRatioMutex.Lock()
+	defer CompletionRatioMutex.Unlock()
 	if CompletionRatio == nil {
 		CompletionRatio = defaultCompletionRatio
 	}
+	return CompletionRatio
+}
+
+func CompletionRatio2JSONString() string {
+	GetCompletionRatioMap()
 	jsonBytes, err := json.Marshal(CompletionRatio)
 	if err != nil {
 		SysError("error marshalling completion ratio: " + err.Error())
@@ -345,12 +356,9 @@ func CompletionRatio2JSONString() string {
 	return string(jsonBytes)
 }
 
-func UpdateCompletionRatioByJSONString(jsonStr string) error {
-	CompletionRatio = make(map[string]float64)
-	return json.Unmarshal([]byte(jsonStr), &CompletionRatio)
-}
-
 func GetCompletionRatio(name string) float64 {
+	GetCompletionRatioMap()
+
 	if strings.Contains(name, "/") {
 		if ratio, ok := CompletionRatio[name]; ok {
 			return ratio
@@ -476,10 +484,3 @@ func GetAudioCompletionRatio(name string) float64 {
 	}
 	return 2
 }
-
-func GetCompletionRatioMap() map[string]float64 {
-	if CompletionRatio == nil {
-		CompletionRatio = defaultCompletionRatio
-	}
-	return CompletionRatio
-}