From d0f76a5c61ded22f4ebb31b8d053d4ad10a24318 Mon Sep 17 00:00:00 2001
From: CalciumIon <1808837298@qq.com>
Date: Fri, 16 Aug 2024 17:25:03 +0800
Subject: [PATCH] feat: support gpt-4o-gizmo-* (close #436)

---
 common/model-ratio.go | 24 ++++++++++++++++--------
 model/cache.go        |  3 +++
 relay/relay-text.go   |  4 ++++
 3 files changed, 23 insertions(+), 8 deletions(-)

diff --git a/common/model-ratio.go b/common/model-ratio.go
index 284fb80..946ee46 100644
--- a/common/model-ratio.go
+++ b/common/model-ratio.go
@@ -23,10 +23,11 @@ const (
 
 var defaultModelRatio = map[string]float64{
 	//"midjourney":                50,
-	"gpt-4-gizmo-*": 15,
-	"gpt-4-all":     15,
-	"gpt-4o-all":    15,
-	"gpt-4":         15,
+	"gpt-4-gizmo-*":  15,
+	"gpt-4o-gizmo-*": 2.5,
+	"gpt-4-all":      15,
+	"gpt-4o-all":     15,
+	"gpt-4":          15,
 	//"gpt-4-0314":                   15, //deprecated
 	"gpt-4-0613": 15,
 	"gpt-4-32k":  30,
@@ -187,8 +188,8 @@ var defaultModelPrice = map[string]float64{
 }
 
 var (
-	modelPriceMap      = make(map[string]float64)
-	modelPriceMapMutex = sync.RWMutex{}
+	modelPriceMap      map[string]float64 = nil
+	modelPriceMapMutex                    = sync.RWMutex{}
 )
 var (
 	modelRatioMap      map[string]float64 = nil
@@ -197,8 +198,9 @@ var (
 
 var CompletionRatio map[string]float64 = nil
 var defaultCompletionRatio = map[string]float64{
-	"gpt-4-gizmo-*": 2,
-	"gpt-4-all":     2,
+	"gpt-4-gizmo-*":  2,
+	"gpt-4o-gizmo-*": 3,
+	"gpt-4-all":      2,
 }
 
 func GetModelPriceMap() map[string]float64 {
@@ -232,6 +234,9 @@ func GetModelPrice(name string, printErr bool) (float64, bool) {
 	if strings.HasPrefix(name, "gpt-4-gizmo") {
 		name = "gpt-4-gizmo-*"
 	}
+	if strings.HasPrefix(name, "gpt-4o-gizmo") {
+		name = "gpt-4o-gizmo-*"
+	}
 	price, ok := modelPriceMap[name]
 	if !ok {
 		if printErr {
@@ -312,6 +317,9 @@ func GetCompletionRatio(name string) float64 {
 	if strings.HasPrefix(name, "gpt-4-gizmo") {
 		name = "gpt-4-gizmo-*"
 	}
+	if strings.HasPrefix(name, "gpt-4o-gizmo") {
+		name = "gpt-4o-gizmo-*"
+	}
 	if strings.HasPrefix(name, "gpt-3.5") {
 		if name == "gpt-3.5-turbo" || strings.HasSuffix(name, "0125") {
 			// https://openai.com/blog/new-embedding-models-and-api-updates
diff --git a/model/cache.go b/model/cache.go
index 2977bb6..7f5411e 100644
--- a/model/cache.go
+++ b/model/cache.go
@@ -270,6 +270,9 @@ func CacheGetRandomSatisfiedChannel(group string, model string, retry int) (*Cha
 	if strings.HasPrefix(model, "gpt-4-gizmo") {
 		model = "gpt-4-gizmo-*"
 	}
+	if strings.HasPrefix(model, "gpt-4o-gizmo") {
+		model = "gpt-4o-gizmo-*"
+	}
 
 	// if memory cache is disabled, get channel directly from database
 	if !common.MemoryCacheEnabled {
diff --git a/relay/relay-text.go b/relay/relay-text.go
index 93d202d..7c6840f 100644
--- a/relay/relay-text.go
+++ b/relay/relay-text.go
@@ -354,6 +354,10 @@ func postConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, modelN
 		logModel = "gpt-4-gizmo-*"
 		logContent += fmt.Sprintf("，模型 %s", modelName)
 	}
+	if strings.HasPrefix(logModel, "gpt-4o-gizmo") {
+		logModel = "gpt-4o-gizmo-*"
+		logContent += fmt.Sprintf("，模型 %s", modelName)
+	}
 	if extraContent != "" {
 		logContent += ", " + extraContent
 	}