From 623fd5f03255f79bdbabbbc2c6ec432b43f4c9dc Mon Sep 17 00:00:00 2001
From: Hao Wu <wuhaoecho@gmail.com>
Date: Sun, 2 Feb 2025 11:30:34 +0800
Subject: [PATCH] Add support for "thought" generation config in Gemini payload

- Introduced `GenerationConfig` struct with `Thought` field in Gemini payload.
- Modified `GenGemminPayload` to accept `thought` parameter and include it in the payload.
- Added logic in `chatStreamGemini` to set `thought` based on the model name.
---
 api/chat_main_handler.go |  6 +++++-
 api/llm/gemini/gemini.go | 16 +++++++++++++++-
 2 files changed, 20 insertions(+), 2 deletions(-)

diff --git a/api/chat_main_handler.go b/api/chat_main_handler.go
index ef606425..bc915636 100644
--- a/api/chat_main_handler.go
+++ b/api/chat_main_handler.go
@@ -1425,7 +1425,11 @@ func (h *ChatHandler) chatStreamGemini(w http.ResponseWriter, chatSession sqlc_q
 		RespondWithError(w, http.StatusInternalServerError, eris.Wrap(err, "Error getting chat files").Error(), err)
 		return "", "", true
 	}
-	payloadBytes, err := gemini.GenGemminPayload(chat_compeletion_messages, chatFiles)
+	thought := false
+	if strings.Contains(chatSession.Model, "thinking") {
+		thought = true
+	}
+	payloadBytes, err := gemini.GenGemminPayload(chat_compeletion_messages, chatFiles, thought)
 	if err != nil {
 		RespondWithError(w, http.StatusInternalServerError, eris.Wrap(err, "Error generating gemmi payload").Error(), err)
 		return "", "", true
diff --git a/api/llm/gemini/gemini.go b/api/llm/gemini/gemini.go
index f3c8f89c..f26ef52e 100644
--- a/api/llm/gemini/gemini.go
+++ b/api/llm/gemini/gemini.go
@@ -67,10 +67,20 @@ func ImageData(mimeType string, data []byte) Blob {
 type GeminiMessage struct {
 	Role  string `json:"role"`
 	Parts []Part `json:"parts"`
+
+}
+
+type GenerationConfig struct {
+	Thought bool `json:"thought"`
+	// Temperature     float64 `json:"temperature"`
+	// TopP            float64 `json:"topP"`
+	// TopK            int     `json:"topK"`
+	// MaxOutputTokens int     `json:"maxOutputTokens"`
 }
 
 type GeminPayload struct {
 	Contents []GeminiMessage `json:"contents"`
+	GenerationConfig  GenerationConfig `json:"generationConfig"`
 }
 
 type Content struct {
@@ -104,6 +114,7 @@ type ResponseBody struct {
 
 func ParseRespLine(line []byte, answer string) string {
 	var resp ResponseBody
+	log.Println(string(line))
 	if err := json.Unmarshal(line, &resp); err != nil {
 		fmt.Println("Failed to parse request body:", err)
 	}
@@ -149,9 +160,12 @@ func SupportedMimeTypes() mapset.Set[string] {
 	)
 }
 
-func GenGemminPayload(chat_compeletion_messages []models.Message, chatFiles []sqlc_queries.ChatFile) ([]byte, error) {
+func GenGemminPayload(chat_compeletion_messages []models.Message, chatFiles []sqlc_queries.ChatFile, thought bool) ([]byte, error) {
 	payload := GeminPayload{
 		Contents: make([]GeminiMessage, len(chat_compeletion_messages)),
+		// GenerationConfig: GenerationConfig{
+		// 	Thought: thought,
+		// },
 	}
 	for i, message := range chat_compeletion_messages {
 		geminiMessage := GeminiMessage{