From f3d5b7766df7dcf0e101f902d04cab30795099ef Mon Sep 17 00:00:00 2001
From: Henry Heng <henryheng@flowiseai.com>
Date: Thu, 30 Oct 2025 22:01:15 +0000
Subject: [PATCH] Feat/add thinking budget to gemini (#5395)

add thinking budget to gemini
---
 .../ChatGoogleGenerativeAI.ts                    | 14 ++++++++++++++
 .../FlowiseChatGoogleGenerativeAI.ts             | 16 ++++++++++++++++
 .../ChatGoogleGenerativeAI/utils/common.ts       |  2 ++
 3 files changed, 32 insertions(+)

diff --git a/packages/components/nodes/chatmodels/ChatGoogleGenerativeAI/ChatGoogleGenerativeAI.ts b/packages/components/nodes/chatmodels/ChatGoogleGenerativeAI/ChatGoogleGenerativeAI.ts
index 95d669e1..d618254c 100644
--- a/packages/components/nodes/chatmodels/ChatGoogleGenerativeAI/ChatGoogleGenerativeAI.ts
+++ b/packages/components/nodes/chatmodels/ChatGoogleGenerativeAI/ChatGoogleGenerativeAI.ts
@@ -174,6 +174,18 @@ class GoogleGenerativeAI_ChatModels implements INode {
                 optional: true,
                 additionalParams: true
             },
+            {
+                label: 'Thinking Budget',
+                name: 'thinkingBudget',
+                type: 'number',
+                description: 'Guides the number of thinking tokens. -1 for dynamic, 0 to disable, or positive integer (Gemini 2.5 models).',
+                step: 1,
+                optional: true,
+                additionalParams: true,
+                show: {
+                    modelName: ['gemini-2.5-pro', 'gemini-2.5-flash', 'gemini-2.5-flash-lite']
+                }
+            },
             {
                 label: 'Base URL',
                 name: 'baseUrl',
@@ -216,6 +228,7 @@ class GoogleGenerativeAI_ChatModels implements INode {
         const cache = nodeData.inputs?.cache as BaseCache
         const streaming = nodeData.inputs?.streaming as boolean
         const baseUrl = nodeData.inputs?.baseUrl as string | undefined
+        const thinkingBudget = nodeData.inputs?.thinkingBudget as string
 
         const allowImageUploads = nodeData.inputs?.allowImageUploads as boolean
 
@@ -235,6 +248,7 @@ class GoogleGenerativeAI_ChatModels implements INode {
         if (cache) obj.cache = cache
         if (temperature) obj.temperature = parseFloat(temperature)
         if (baseUrl) obj.baseUrl = baseUrl
+        if (thinkingBudget) obj.thinkingBudget = parseInt(thinkingBudget, 10)
 
         let safetySettings: SafetySetting[] = []
         if (_safetySettings) {
diff --git a/packages/components/nodes/chatmodels/ChatGoogleGenerativeAI/FlowiseChatGoogleGenerativeAI.ts b/packages/components/nodes/chatmodels/ChatGoogleGenerativeAI/FlowiseChatGoogleGenerativeAI.ts
index 5c21fa4c..cdf3ac11 100644
--- a/packages/components/nodes/chatmodels/ChatGoogleGenerativeAI/FlowiseChatGoogleGenerativeAI.ts
+++ b/packages/components/nodes/chatmodels/ChatGoogleGenerativeAI/FlowiseChatGoogleGenerativeAI.ts
@@ -174,6 +174,9 @@ export interface GoogleGenerativeAIChatInput extends BaseChatModelParams, Pick<G
      * - Gemini 1.0 Pro version gemini-1.0-pro-002
      */
     convertSystemMessageToHumanContent?: boolean | undefined
+
+    /** Thinking budget for Gemini 2.5 thinking models. Supports -1 (dynamic), 0 (off), or positive integers. */
+    thinkingBudget?: number
 }
 
 /**
@@ -599,6 +602,8 @@ export class LangchainChatGoogleGenerativeAI
 
     convertSystemMessageToHumanContent: boolean | undefined
 
+    thinkingBudget?: number
+
     private client: GenerativeModel
 
     get _isMultimodalModel() {
@@ -657,6 +662,7 @@ export class LangchainChatGoogleGenerativeAI
 
         this.streaming = fields.streaming ?? this.streaming
         this.json = fields.json
+        this.thinkingBudget = fields.thinkingBudget
 
         this.client = new GenerativeAI(this.apiKey).getGenerativeModel(
             {
@@ -676,12 +682,22 @@ export class LangchainChatGoogleGenerativeAI
                 baseUrl: fields.baseUrl
             }
         )
+        if (this.thinkingBudget !== undefined) {
+            ;(this.client.generationConfig as any).thinkingConfig = {
+                ...(this.thinkingBudget !== undefined ? { thinkingBudget: this.thinkingBudget } : {})
+            }
+        }
         this.streamUsage = fields.streamUsage ?? this.streamUsage
     }
 
     useCachedContent(cachedContent: CachedContent, modelParams?: ModelParams, requestOptions?: RequestOptions): void {
         if (!this.apiKey) return
         this.client = new GenerativeAI(this.apiKey).getGenerativeModelFromCachedContent(cachedContent, modelParams, requestOptions)
+        if (this.thinkingBudget !== undefined) {
+            ;(this.client.generationConfig as any).thinkingConfig = {
+                ...(this.thinkingBudget !== undefined ? { thinkingBudget: this.thinkingBudget } : {})
+            }
+        }
     }
 
     get useSystemInstruction(): boolean {
diff --git a/packages/components/nodes/chatmodels/ChatGoogleGenerativeAI/utils/common.ts b/packages/components/nodes/chatmodels/ChatGoogleGenerativeAI/utils/common.ts
index 3799ba3b..13d40e18 100644
--- a/packages/components/nodes/chatmodels/ChatGoogleGenerativeAI/utils/common.ts
+++ b/packages/components/nodes/chatmodels/ChatGoogleGenerativeAI/utils/common.ts
@@ -48,6 +48,8 @@ export function getMessageAuthor(message: BaseMessage) {
 }
 
 /**
+ * !!! IMPORANT: Must return 'user' as default instead of throwing error
+ * https://github.com/FlowiseAI/Flowise/issues/4743
  * Maps a message type to a Google Generative AI chat author.
  * @param message The message to map.
  * @param model The model to use for mapping.