From 94f3223fd74829555e831b5a011c6a88da25b7f9 Mon Sep 17 00:00:00 2001
From: RandyZhang <zhangshi3421@gmail.com>
Date: Thu, 15 May 2025 16:00:42 +0800
Subject: [PATCH] =?UTF-8?q?=E6=94=AF=E6=8C=81vertexai=20gemini=E7=9A=84thi?=
 =?UTF-8?q?nking=E6=8E=A7=E5=88=B6?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 relay/adaptor/gemini/model.go            | 24 ++++++++------
 relay/adaptor/vertexai/gemini/adapter.go | 41 ++++++++++++++++++++++--
 2 files changed, 54 insertions(+), 11 deletions(-)

diff --git a/relay/adaptor/gemini/model.go b/relay/adaptor/gemini/model.go
index c3acae60..6469e45f 100644
--- a/relay/adaptor/gemini/model.go
+++ b/relay/adaptor/gemini/model.go
@@ -65,13 +65,19 @@ type ChatTools struct {
 	FunctionDeclarations any `json:"function_declarations,omitempty"`
 }
 
-type ChatGenerationConfig struct {
-	ResponseMimeType string   `json:"responseMimeType,omitempty"`
-	ResponseSchema   any      `json:"responseSchema,omitempty"`
-	Temperature      *float64 `json:"temperature,omitempty"`
-	TopP             *float64 `json:"topP,omitempty"`
-	TopK             float64  `json:"topK,omitempty"`
-	MaxOutputTokens  int      `json:"maxOutputTokens,omitempty"`
-	CandidateCount   int      `json:"candidateCount,omitempty"`
-	StopSequences    []string `json:"stopSequences,omitempty"`
+type ThinkingConfig struct {
+	IncludeThoughts bool `json:"includeThoughts"`
+	ThinkingBudget  int  `json:"thinkingBudget"`
+}
+
+type ChatGenerationConfig struct {
+	ResponseMimeType string          `json:"responseMimeType,omitempty"`
+	ResponseSchema   any             `json:"responseSchema,omitempty"`
+	Temperature      *float64        `json:"temperature,omitempty"`
+	TopP             *float64        `json:"topP,omitempty"`
+	TopK             float64         `json:"topK,omitempty"`
+	MaxOutputTokens  int             `json:"maxOutputTokens,omitempty"`
+	CandidateCount   int             `json:"candidateCount,omitempty"`
+	StopSequences    []string        `json:"stopSequences,omitempty"`
+	ThinkingConfig   *ThinkingConfig `json:"thinkingConfig,omitempty"`
 }
diff --git a/relay/adaptor/vertexai/gemini/adapter.go b/relay/adaptor/vertexai/gemini/adapter.go
index f5b245d8..abdd13c7 100644
--- a/relay/adaptor/vertexai/gemini/adapter.go
+++ b/relay/adaptor/vertexai/gemini/adapter.go
@@ -2,6 +2,9 @@ package vertexai
 
 import (
 	"net/http"
+	"net/url"
+	"strconv"
+	"strings"
 
 	"github.com/gin-gonic/gin"
 	"github.com/pkg/errors"
@@ -27,13 +30,47 @@ var ModelList = []string{
 type Adaptor struct {
 }
 
+func (a *Adaptor) parseGeminiChatGenerationThinking(model string) (string, *gemini.ThinkingConfig) {
+	thinkingConfig := &gemini.ThinkingConfig{
+		IncludeThoughts: false,
+		ThinkingBudget:  0,
+	}
+	modelName := model
+	if strings.Contains(model, "?") {
+		parts := strings.Split(model, "?")
+		_modelName := parts[0]
+		if len(parts) >= 2 {
+			modelOptions, err := url.ParseQuery(parts[1])
+			if err != nil && modelOptions != nil {
+				modelName = _modelName
+				enableThinking := modelOptions.Has("thinking")
+				if enableThinking {
+					thinkingConfig.IncludeThoughts = true
+				}
+				thinkingBudget := modelOptions.Get("thinking_budget")
+				if thinkingBudget != "" {
+					thinkingBudgetInt, err := strconv.Atoi(thinkingBudget)
+					if err != nil {
+						thinkingConfig.ThinkingBudget = thinkingBudgetInt
+					}
+				}
+			}
+		}
+	}
+	return modelName, thinkingConfig
+}
+
 func (a *Adaptor) ConvertRequest(c *gin.Context, relayMode int, request *model.GeneralOpenAIRequest) (any, error) {
 	if request == nil {
 		return nil, errors.New("request is nil")
 	}
-
+	modelName, thinkingConfig := a.parseGeminiChatGenerationThinking(request.Model)
+	request.Model = modelName
 	geminiRequest := gemini.ConvertRequest(*request)
-	c.Set(ctxkey.RequestModel, request.Model)
+	if thinkingConfig != nil {
+		geminiRequest.GenerationConfig.ThinkingConfig = thinkingConfig
+	}
+	c.Set(ctxkey.RequestModel, modelName)
 	c.Set(ctxkey.ConvertedRequest, geminiRequest)
 	return geminiRequest, nil
 }