feat(go/plugins/googlegenai): add Thinking support

go/plugins/googlegenai/gemini.go

@@ -39,6 +39,14 @@ import (
 	"google.golang.org/genai"
 )
 
+const (
+	// Thinking budget limit
+	thinkingBudgetMax = 24576
+
+	// Tool name regex
+	toolNameRegex = "^[a-zA-Z_][a-zA-Z0-9_.-]{0,63}$"
+)
+
 var (
 	// BasicText describes model capabilities for text-only Gemini models.
 	BasicText = ai.ModelSupports{
@@ -59,9 +67,6 @@ var (
 		Constrained: ai.ConstrainedSupportNoTools,
 	}
 
-	// Tool name regex
-	toolNameRegex = "^[a-zA-Z_][a-zA-Z0-9_.-]{0,63}$"
-
 	// Attribution header
 	xGoogApiClientHeader = http.CanonicalHeaderKey("x-goog-api-client")
 	genkitClientHeader   = http.Header{
@@ -175,6 +180,14 @@ type SafetySetting struct {
 	Threshold HarmBlockThreshold `json:"threshold,omitempty"`
 }
 
+// Thinking configuration to control reasoning
+type ThinkingConfig struct {
+	// Indicates whether the response should include thoughts (if available and supported)
+	IncludeThoughts bool `json:"includeThoughts,omitempty"`
+	// Thinking budget in tokens. If set to zero, thinking gets disabled
+	ThinkingBudget int32 `json:"thinkingBudget,omitempty"`
+}
+
 type Modality string
 
 const (
@@ -204,6 +217,8 @@ type GeminiConfig struct {
 	CodeExecution bool `json:"codeExecution,omitempty"`
 	// Response modalities for returned model messages
 	ResponseModalities []Modality `json:"responseModalities,omitempty"`
+	// Thinking configuration controls the model's internal reasoning process
+	ThinkingConfig *ThinkingConfig `json:"thinkingConfig,omitempty"`
 }
 
 // configFromRequest converts any supported config type to [GeminiConfig].
@@ -390,6 +405,7 @@ func generate(
 			return nil, err
 		}
 		r := translateResponse(resp)
+
 		r.Request = input
 		if cache != nil {
 			r.Message.Metadata = setCacheMetadata(r.Message.Metadata, cache)
@@ -528,6 +544,16 @@ func toGeminiRequest(input *ai.ModelRequest, cache *genai.CachedContent) (*genai
 		})
 	}
 
+	if c.ThinkingConfig != nil {
+		if c.ThinkingConfig.ThinkingBudget < 0 || c.ThinkingConfig.ThinkingBudget > thinkingBudgetMax {
+			return nil, fmt.Errorf("thinkingBudget should be between 0 and %d", thinkingBudgetMax)
+		}
+		gcc.ThinkingConfig = &genai.ThinkingConfig{
+			IncludeThoughts: c.ThinkingConfig.IncludeThoughts,
+			ThinkingBudget:  &c.ThinkingConfig.ThinkingBudget,
+		}
+	}
+
 	var systemParts []*genai.Part
 	for _, m := range input.Messages {
 		if m.Role == ai.RoleSystem {
@@ -767,6 +793,10 @@ func translateCandidate(cand *genai.Candidate) *ai.ModelResponse {
 
 		if part.Text != "" {
 			partFound++
+			if part.Thought {
+				// TODO: Include a `reasoning` part. Not available in the SDK yet.
+				continue
+			}
 			p = ai.NewTextPart(part.Text)
 		}
 		if part.InlineData != nil {

go/plugins/googlegenai/gemini_test.go

@@ -44,6 +44,10 @@ func TestConvertRequest(t *testing.T) {
 			TopK:            1.0,
 			TopP:            1.0,
 			Version:         text,
+			ThinkingConfig: &ThinkingConfig{
+				IncludeThoughts: false,
+				ThinkingBudget:  0,
+			},
 		},
 		Tools:      []*ai.ToolDefinition{tool},
 		ToolChoice: ai.ToolChoiceAuto,
@@ -143,6 +147,30 @@ func TestConvertRequest(t *testing.T) {
 		if gcc.ResponseSchema == nil {
 			t.Errorf("ResponseSchema should not be empty")
 		}
+		if gcc.ThinkingConfig == nil {
+			t.Errorf("ThinkingConfig should not be empty")
+		}
+	})
+	t.Run("thinking budget limits", func(t *testing.T) {
+		thinkingBudget := GeminiConfig{
+			ThinkingConfig: &ThinkingConfig{
+				IncludeThoughts: false,
+				ThinkingBudget:  -23,
+			},
+		}
+		req := &ai.ModelRequest{
+			Config: thinkingBudget,
+		}
+		_, err := toGeminiRequest(req, nil)
+		if err == nil {
+			t.Fatal("expecting an error, thinking budget should not be negative")
+		}
+		thinkingBudget.ThinkingConfig.ThinkingBudget = 999999
+		req.Config = thinkingBudget
+		_, err = toGeminiRequest(req, nil)
+		if err == nil {
+			t.Fatalf("expecting an error, thinking budget should not be greater than %d", thinkingBudgetMax)
+		}
 	})
 	t.Run("convert tools with valid tool", func(t *testing.T) {
 		tools := []*ai.ToolDefinition{tool}

go/plugins/googlegenai/googleai_live_test.go

@@ -401,6 +401,50 @@ func TestGoogleAILive(t *testing.T) {
 			t.Errorf("Empty usage stats %#v", *resp.Usage)
 		}
 	})
+	t.Run("thinking", func(t *testing.T) {
+		m := googlegenai.GoogleAIModel(g, "gemini-2.5-flash-preview-04-17")
+		resp, err := genkit.Generate(ctx, g,
+			ai.WithConfig(googlegenai.GeminiConfig{
+				Temperature: 0.4,
+				ThinkingConfig: &googlegenai.ThinkingConfig{
+					IncludeThoughts: true,
+					ThinkingBudget:  100,
+				},
+			}),
+			ai.WithModel(m),
+			ai.WithPrompt("Analogize photosynthesis and growing up."))
+		if err != nil {
+			t.Fatal(err)
+		}
+		if resp == nil {
+			t.Fatal("nil response obtanied")
+		}
+		if resp.Usage.ThoughtsTokens == 0 || resp.Usage.ThoughtsTokens > 100 {
+			t.Fatal("thoughts tokens should not be zero or greater than 100")
+		}
+	})
+	t.Run("thinking disabled", func(t *testing.T) {
+		m := googlegenai.GoogleAIModel(g, "gemini-2.5-flash-preview-04-17")
+		resp, err := genkit.Generate(ctx, g,
+			ai.WithConfig(googlegenai.GeminiConfig{
+				Temperature: 0.4,
+				ThinkingConfig: &googlegenai.ThinkingConfig{
+					IncludeThoughts: false,
+					ThinkingBudget:  0,
+				},
+			}),
+			ai.WithModel(m),
+			ai.WithPrompt("Analogize photosynthesis and growing up."))
+		if err != nil {
+			t.Fatal(err)
+		}
+		if resp == nil {
+			t.Fatal("nil response obtanied")
+		}
+		if resp.Usage.ThoughtsTokens > 0 {
+			t.Fatal("thoughts tokens should be zero")
+		}
+	})
 }
 
 func TestCacheHelper(t *testing.T) {