From 49ffb1c60dcfed8b361f8aa6202e190a0239a38f Mon Sep 17 00:00:00 2001
From: Wei Tingjiang <mxdlzg@163.com>
Date: Sun, 22 Dec 2024 18:25:44 +0800
Subject: [PATCH] feat: enhance response handling to support
 gemini-2.0-thinking (#1995)

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

diff --git a/relay/adaptor/gemini/constants.go b/relay/adaptor/gemini/constants.go
index fa53d63baf..9d1cbc4acd 100644
--- a/relay/adaptor/gemini/constants.go
+++ b/relay/adaptor/gemini/constants.go
@@ -7,4 +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",
 }
diff --git a/relay/adaptor/gemini/main.go b/relay/adaptor/gemini/main.go
index d6ab45d489..3aa4974b83 100644
--- a/relay/adaptor/gemini/main.go
+++ b/relay/adaptor/gemini/main.go
@@ -55,6 +55,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,
@@ -247,7 +251,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 0557b075c6..b537787553 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 4a7793df6d..8fa708f021 100644
--- a/relay/billing/ratio/model.go
+++ b/relay/billing/ratio/model.go
@@ -108,14 +108,15 @@ 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,
-	"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,
+	"aqa":                           1,
 	// https://open.bigmodel.cn/pricing
 	"glm-4":         0.1 * RMB,
 	"glm-4v":        0.1 * RMB,