From 8cd1db4e9964798c13bec62a3e7c2ee75eeb88d1 Mon Sep 17 00:00:00 2001
From: "Laisky.Cai" <github@laisky.com>
Date: Thu, 19 Dec 2024 05:55:21 +0000
Subject: [PATCH] feat: add support for new OpenAI models and update billing
 ratios

---
 relay/adaptor/openai/constants.go |  3 +++
 relay/billing/ratio/model.go      | 14 ++++++++++++--
 2 files changed, 15 insertions(+), 2 deletions(-)

diff --git a/relay/adaptor/openai/constants.go b/relay/adaptor/openai/constants.go
index aacdba1a..be4804c2 100644
--- a/relay/adaptor/openai/constants.go
+++ b/relay/adaptor/openai/constants.go
@@ -20,4 +20,7 @@ var ModelList = []string{
 	"dall-e-2", "dall-e-3",
 	"whisper-1",
 	"tts-1", "tts-1-1106", "tts-1-hd", "tts-1-hd-1106",
+	"o1", "o1-2024-12-17",
+	"o1-preview", "o1-preview-2024-09-12",
+	"o1-mini", "o1-mini-2024-09-12",
 }
diff --git a/relay/billing/ratio/model.go b/relay/billing/ratio/model.go
index 1b58ec09..96285126 100644
--- a/relay/billing/ratio/model.go
+++ b/relay/billing/ratio/model.go
@@ -48,8 +48,14 @@ var ModelRatio = map[string]float64{
 	"gpt-3.5-turbo-instruct":  0.75, // $0.0015 / 1K tokens
 	"gpt-3.5-turbo-1106":      0.5,  // $0.001 / 1K tokens
 	"gpt-3.5-turbo-0125":      0.25, // $0.0005 / 1K tokens
-	"davinci-002":             1,    // $0.002 / 1K tokens
-	"babbage-002":             0.2,  // $0.0004 / 1K tokens
+	"o1":                      7.5,  // $15.00 / 1M input tokens
+	"o1-2024-12-17":           7.5,
+	"o1-preview":              7.5, // $15.00 / 1M input tokens
+	"o1-preview-2024-09-12":   7.5,
+	"o1-mini":                 1.5, // $3.00 / 1M input tokens
+	"o1-mini-2024-09-12":      1.5,
+	"davinci-002":             1,   // $0.002 / 1K tokens
+	"babbage-002":             0.2, // $0.0004 / 1K tokens
 	"text-ada-001":            0.2,
 	"text-babbage-001":        0.25,
 	"text-curie-001":          1,
@@ -344,6 +350,10 @@ func GetCompletionRatio(name string, channelType int) float64 {
 		}
 		return 2
 	}
+	// including o1, o1-preview, o1-mini
+	if strings.HasPrefix(name, "o1") {
+		return 4
+	}
 	if name == "chatgpt-4o-latest" {
 		return 3
 	}