From 32ea6676c7e24f2271b5f69cf141253e6df963fc Mon Sep 17 00:00:00 2001
From: "Laisky.Cai" <github@laisky.com>
Date: Wed, 22 Jan 2025 03:27:22 +0000
Subject: [PATCH] feat: add support for new Gemini model version
 'gemini-2.0-flash-thinking-exp-01-21'

---
 relay/adaptor/gemini/adaptor.go          |  8 +++++---
 relay/adaptor/gemini/constants.go        |  2 +-
 relay/adaptor/vertexai/gemini/adapter.go |  3 ++-
 relay/billing/ratio/model.go             | 19 ++++++++++---------
 4 files changed, 18 insertions(+), 14 deletions(-)

diff --git a/relay/adaptor/gemini/adaptor.go b/relay/adaptor/gemini/adaptor.go
index a86fde40..edae1791 100644
--- a/relay/adaptor/gemini/adaptor.go
+++ b/relay/adaptor/gemini/adaptor.go
@@ -7,7 +7,6 @@ import (
 	"net/http"
 
 	"github.com/gin-gonic/gin"
-	"github.com/songquanpeng/one-api/common/config"
 	"github.com/songquanpeng/one-api/common/helper"
 	channelhelper "github.com/songquanpeng/one-api/relay/adaptor"
 	"github.com/songquanpeng/one-api/relay/adaptor/openai"
@@ -24,8 +23,11 @@ func (a *Adaptor) Init(meta *meta.Meta) {
 }
 
 func (a *Adaptor) GetRequestURL(meta *meta.Meta) (string, error) {
-	defaultVersion := config.GeminiVersion
-	if meta.ActualModelName == "gemini-2.0-flash-exp" {
+	var defaultVersion string
+	switch meta.ActualModelName {
+	case "gemini-2.0-flash-exp",
+		"gemini-2.0-flash-thinking-exp",
+		"gemini-2.0-flash-thinking-exp-01-21":
 		defaultVersion = "v1beta"
 	}
 
diff --git a/relay/adaptor/gemini/constants.go b/relay/adaptor/gemini/constants.go
index 9d1cbc4a..381d0c12 100644
--- a/relay/adaptor/gemini/constants.go
+++ b/relay/adaptor/gemini/constants.go
@@ -7,5 +7,5 @@ var ModelList = []string{
 	"gemini-1.5-flash", "gemini-1.5-pro",
 	"text-embedding-004", "aqa",
 	"gemini-2.0-flash-exp",
-	"gemini-2.0-flash-thinking-exp",
+	"gemini-2.0-flash-thinking-exp", "gemini-2.0-flash-thinking-exp-01-21",
 }
diff --git a/relay/adaptor/vertexai/gemini/adapter.go b/relay/adaptor/vertexai/gemini/adapter.go
index b5377875..1240ea5b 100644
--- a/relay/adaptor/vertexai/gemini/adapter.go
+++ b/relay/adaptor/vertexai/gemini/adapter.go
@@ -18,7 +18,8 @@ var ModelList = []string{
 	"gemini-pro", "gemini-pro-vision",
 	"gemini-1.5-pro-001", "gemini-1.5-flash-001",
 	"gemini-1.5-pro-002", "gemini-1.5-flash-002",
-	"gemini-2.0-flash-exp", "gemini-2.0-flash-thinking-exp",
+	"gemini-2.0-flash-exp",
+	"gemini-2.0-flash-thinking-exp", "gemini-2.0-flash-thinking-exp-01-21",
 }
 
 type Adaptor struct {
diff --git a/relay/billing/ratio/model.go b/relay/billing/ratio/model.go
index 71ca0e95..7fe08506 100644
--- a/relay/billing/ratio/model.go
+++ b/relay/billing/ratio/model.go
@@ -110,15 +110,16 @@ var ModelRatio = map[string]float64{
 	"bge-large-en":       0.002 * RMB,
 	"tao-8k":             0.002 * RMB,
 	// https://ai.google.dev/pricing
-	"gemini-pro":                    1, // $0.00025 / 1k characters -> $0.001 / 1k tokens
-	"gemini-1.0-pro":                1,
-	"gemini-1.5-pro":                1,
-	"gemini-1.5-pro-001":            1,
-	"gemini-1.5-flash":              1,
-	"gemini-1.5-flash-001":          1,
-	"gemini-2.0-flash-exp":          1,
-	"gemini-2.0-flash-thinking-exp": 1,
-	"aqa":                           1,
+	"gemini-pro":                          1, // $0.00025 / 1k characters -> $0.001 / 1k tokens
+	"gemini-1.0-pro":                      1,
+	"gemini-1.5-pro":                      1,
+	"gemini-1.5-pro-001":                  1,
+	"gemini-1.5-flash":                    1,
+	"gemini-1.5-flash-001":                1,
+	"gemini-2.0-flash-exp":                1,
+	"gemini-2.0-flash-thinking-exp":       1,
+	"gemini-2.0-flash-thinking-exp-01-21": 1,
+	"aqa":                                 1,
 	// https://open.bigmodel.cn/pricing
 	"glm-4":         0.1 * RMB,
 	"glm-4v":        0.1 * RMB,