From c4afcc8337188559326eae339a35767d761a94da Mon Sep 17 00:00:00 2001
From: Jinjun Liu <ljj7129@gmail.com>
Date: Sat, 21 Dec 2024 19:46:35 -0600
Subject: [PATCH] feat: add support for gemini-2.0-flash-thinking-exp model

---
 relay/adaptor/gemini/adaptor.go          |  3 ++-
 relay/adaptor/gemini/constants.go        |  2 +-
 relay/adaptor/gemini/main.go             | 13 ++++++++++++-
 relay/adaptor/vertexai/gemini/adapter.go |  2 +-
 relay/billing/ratio/model.go             |  1 +
 5 files changed, 17 insertions(+), 4 deletions(-)

diff --git a/relay/adaptor/gemini/adaptor.go b/relay/adaptor/gemini/adaptor.go
index 4178031e..f7df87a6 100644
--- a/relay/adaptor/gemini/adaptor.go
+++ b/relay/adaptor/gemini/adaptor.go
@@ -25,7 +25,8 @@ 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" {
+	// gemini-2.0-flash-exp and gemini-2.0-flash-thinking-exp use v1beta
+	if meta.ActualModelName == "gemini-2.0-flash-exp" || meta.ActualModelName == "gemini-2.0-flash-thinking-exp" {
 		defaultVersion = "v1beta"
 	}
 
diff --git a/relay/adaptor/gemini/constants.go b/relay/adaptor/gemini/constants.go
index fa53d63b..57b14378 100644
--- a/relay/adaptor/gemini/constants.go
+++ b/relay/adaptor/gemini/constants.go
@@ -6,5 +6,5 @@ var ModelList = []string{
 	"gemini-pro", "gemini-1.0-pro",
 	"gemini-1.5-flash", "gemini-1.5-pro",
 	"text-embedding-004", "aqa",
-	"gemini-2.0-flash-exp",
+	"gemini-2.0-flash-exp", "gemini-2.0-flash-thinking-exp",
 }
diff --git a/relay/adaptor/gemini/main.go b/relay/adaptor/gemini/main.go
index 86832e53..eca9b010 100644
--- a/relay/adaptor/gemini/main.go
+++ b/relay/adaptor/gemini/main.go
@@ -54,6 +54,10 @@ func ConvertRequest(textRequest model.GeneralOpenAIRequest) *ChatRequest {
 				Category:  "HARM_CATEGORY_DANGEROUS_CONTENT",
 				Threshold: config.GeminiSafetySetting,
 			},
+			{
+				Category:  "HARM_CATEGORY_CIVIC_INTEGRITY",
+				Threshold: config.GeminiSafetySetting,
+			},
 		},
 		GenerationConfig: ChatGenerationConfig{
 			Temperature:     textRequest.Temperature,
@@ -246,7 +250,14 @@ func responseGeminiChat2OpenAI(response *ChatResponse) *openai.TextResponse {
 			if candidate.Content.Parts[0].FunctionCall != nil {
 				choice.Message.ToolCalls = getToolCalls(&candidate)
 			} else {
-				choice.Message.Content = candidate.Content.Parts[0].Text
+				var builder strings.Builder
+				for _, part := range candidate.Content.Parts {
+					if i > 0 {
+						builder.WriteString("\n")
+					}
+					builder.WriteString(part.Text)
+				}
+				choice.Message.Content = builder.String()
 			}
 		} else {
 			choice.Message.Content = ""
diff --git a/relay/adaptor/vertexai/gemini/adapter.go b/relay/adaptor/vertexai/gemini/adapter.go
index 0557b075..b5377875 100644
--- a/relay/adaptor/vertexai/gemini/adapter.go
+++ b/relay/adaptor/vertexai/gemini/adapter.go
@@ -18,7 +18,7 @@ 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-exp", "gemini-2.0-flash-thinking-exp",
 }
 
 type Adaptor struct {
diff --git a/relay/billing/ratio/model.go b/relay/billing/ratio/model.go
index f58219ef..3b867867 100644
--- a/relay/billing/ratio/model.go
+++ b/relay/billing/ratio/model.go
@@ -117,6 +117,7 @@ var ModelRatio = map[string]float64{
 	"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,
 	// https://open.bigmodel.cn/pricing
 	"glm-4":         0.1 * RMB,