feat: enhance response handling to support gemini-2.0-thinking

songquanpeng · Dec 21, 2024 · 89eb083 · 89eb083
1 parent e9f117f
commit 89eb083
Show file tree

Hide file tree

Showing 4 changed files with 23 additions and 10 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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,