From 3af62a3efabad767c5724bab456904cd33bf5fcd Mon Sep 17 00:00:00 2001
From: CalciumIon <1808837298@qq.com>
Date: Fri, 13 Sep 2024 01:22:27 +0800
Subject: [PATCH] feat: support OpenAI o1-preview and o1-mini

---
 common/model-ratio.go            | 29 ++++++++++++++++++-----------
 relay/channel/openai/constant.go |  2 ++
 2 files changed, 20 insertions(+), 11 deletions(-)

diff --git a/common/model-ratio.go b/common/model-ratio.go
index 4916449..b45ed29 100644
--- a/common/model-ratio.go
+++ b/common/model-ratio.go
@@ -42,6 +42,10 @@ var defaultModelRatio = map[string]float64{
 	"gpt-4o":                    2.5,  // $0.01 / 1K tokens
 	"gpt-4o-2024-05-13":         2.5,  // $0.01 / 1K tokens
 	"gpt-4o-2024-08-06":         1.25, // $0.01 / 1K tokens
+	"o1-preview":                7.5,
+	"o1-preview-2024-09-12":     7.5,
+	"o1-mini":                   1.5,
+	"o1-mini-2024-09-12":        1.5,
 	"gpt-4o-mini":               0.075,
 	"gpt-4o-mini-2024-07-18":    0.075,
 	"gpt-4-turbo":               5,    // $0.01 / 1K tokens
@@ -331,17 +335,6 @@ func GetCompletionRatio(name string) float64 {
 	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
-			// Updated GPT-3.5 Turbo model and lower pricing
-			return 3
-		}
-		if strings.HasSuffix(name, "1106") {
-			return 2
-		}
-		return 4.0 / 3.0
-	}
 	if strings.HasPrefix(name, "gpt-4") && !strings.HasSuffix(name, "-all") && !strings.HasSuffix(name, "-gizmo-*") {
 		if strings.HasPrefix(name, "gpt-4-turbo") || strings.HasSuffix(name, "preview") {
 			return 3
@@ -354,6 +347,9 @@ func GetCompletionRatio(name string) float64 {
 		}
 		return 2
 	}
+	if strings.HasPrefix(name, "o1-") {
+		return 4
+	}
 	if name == "chatgpt-4o-latest" {
 		return 3
 	}
@@ -364,6 +360,17 @@ func GetCompletionRatio(name string) float64 {
 	} else if strings.Contains(name, "claude-3") {
 		return 5
 	}
+	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
+			// Updated GPT-3.5 Turbo model and lower pricing
+			return 3
+		}
+		if strings.HasSuffix(name, "1106") {
+			return 2
+		}
+		return 4.0 / 3.0
+	}
 	if strings.HasPrefix(name, "mistral-") {
 		return 3
 	}
diff --git a/relay/channel/openai/constant.go b/relay/channel/openai/constant.go
index ac2d673..5e5fd21 100644
--- a/relay/channel/openai/constant.go
+++ b/relay/channel/openai/constant.go
@@ -11,6 +11,8 @@ var ModelList = []string{
 	"chatgpt-4o-latest",
 	"gpt-4o", "gpt-4o-2024-05-13", "gpt-4o-2024-08-06",
 	"gpt-4o-mini", "gpt-4o-mini-2024-07-18",
+	"o1-preview", "o1-preview-2024-09-12",
+	"o1-mini", "o1-mini-2024-09-12",
 	"text-embedding-ada-002", "text-embedding-3-small", "text-embedding-3-large",
 	"text-curie-001", "text-babbage-001", "text-ada-001",
 	"text-moderation-latest", "text-moderation-stable",