From 9066b3e12d60e8bc2d34372bf0e7cd87c3e96466 Mon Sep 17 00:00:00 2001 From: Hugo Aguirre Parra Date: Mon, 28 Apr 2025 20:20:30 +0000 Subject: [PATCH 1/4] initial commit for thinking configuration (no parts) --- go/plugins/googlegenai/gemini.go | 20 ++++++++++++++++++ go/plugins/googlegenai/gemini_test.go | 7 +++++++ go/plugins/googlegenai/googleai_live_test.go | 22 ++++++++++++++++++++ 3 files changed, 49 insertions(+) diff --git a/go/plugins/googlegenai/gemini.go b/go/plugins/googlegenai/gemini.go index d866763ec..135a31a4c 100644 --- a/go/plugins/googlegenai/gemini.go +++ b/go/plugins/googlegenai/gemini.go @@ -175,6 +175,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 +212,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]. @@ -528,6 +538,13 @@ func toGeminiRequest(input *ai.ModelRequest, cache *genai.CachedContent) (*genai }) } + if c.ThinkingConfig != nil { + 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 { @@ -798,6 +815,9 @@ func translateCandidate(cand *genai.Candidate) *ai.ModelResponse { part.ExecutableCode.Code, ) } + if part.Thought { + fmt.Printf("THOUGHT found!!\n\n") + } if partFound > 1 { panic(fmt.Sprintf("expected only 1 content part in response, got %d, part: %#v", partFound, part)) } diff --git a/go/plugins/googlegenai/gemini_test.go b/go/plugins/googlegenai/gemini_test.go index 2e81d4aae..8d8244fc8 100644 --- a/go/plugins/googlegenai/gemini_test.go +++ b/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, @@ -121,6 +125,9 @@ 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("convert tools with valid tool", func(t *testing.T) { tools := []*ai.ToolDefinition{tool} diff --git a/go/plugins/googlegenai/googleai_live_test.go b/go/plugins/googlegenai/googleai_live_test.go index 61a844b96..4e9e1f0d6 100644 --- a/go/plugins/googlegenai/googleai_live_test.go +++ b/go/plugins/googlegenai/googleai_live_test.go @@ -398,6 +398,28 @@ 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: 200, + }, + }), + ai.WithModel(m), + ai.WithPrompt("Analogize photosynthesis and growing up.")) + if err != nil { + t.Fatal(err) + } + if resp == nil { + t.Fatal("nil response obtanied") + } + // since Thinking was enabled, the response should have thinking parts + fmt.Print(resp.Text()) + t.Fatal("remove me") + }) } func TestCacheHelper(t *testing.T) { From f4834ed0fc12bf669de4f3fb2b7190ba3f00b00d Mon Sep 17 00:00:00 2001 From: Hugo Aguirre Parra Date: Wed, 30 Apr 2025 20:04:43 +0000 Subject: [PATCH 2/4] minor changes and debug logs --- go/plugins/googlegenai/gemini.go | 5 +++++ go/plugins/googlegenai/googleai_live_test.go | 2 +- 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/go/plugins/googlegenai/gemini.go b/go/plugins/googlegenai/gemini.go index 135a31a4c..4e0f38acb 100644 --- a/go/plugins/googlegenai/gemini.go +++ b/go/plugins/googlegenai/gemini.go @@ -399,6 +399,7 @@ func generate( if err != nil { return nil, err } + fmt.Printf("thinking tokens: %d", resp.UsageMetadata.ThoughtsTokenCount) r := translateResponse(resp) r.Request = input if cache != nil { @@ -539,6 +540,7 @@ func toGeminiRequest(input *ai.ModelRequest, cache *genai.CachedContent) (*genai } if c.ThinkingConfig != nil { + fmt.Printf("setting thinking config!\n\n") gcc.ThinkingConfig = &genai.ThinkingConfig{ IncludeThoughts: c.ThinkingConfig.IncludeThoughts, ThinkingBudget: &c.ThinkingConfig.ThinkingBudget, @@ -785,6 +787,9 @@ func translateCandidate(cand *genai.Candidate) *ai.ModelResponse { if part.Text != "" { partFound++ p = ai.NewTextPart(part.Text) + if part.Thought { + fmt.Printf("this is a text part, and is a thought!\n\n part: %q", part.Text) + } } if part.InlineData != nil { partFound++ diff --git a/go/plugins/googlegenai/googleai_live_test.go b/go/plugins/googlegenai/googleai_live_test.go index 4e9e1f0d6..6b346e8cc 100644 --- a/go/plugins/googlegenai/googleai_live_test.go +++ b/go/plugins/googlegenai/googleai_live_test.go @@ -405,7 +405,7 @@ func TestGoogleAILive(t *testing.T) { Temperature: 0.4, ThinkingConfig: &googlegenai.ThinkingConfig{ IncludeThoughts: true, - ThinkingBudget: 200, + ThinkingBudget: 1024, }, }), ai.WithModel(m), From a316887fb947a63aaea2199413dbe36bf4cf0d45 Mon Sep 17 00:00:00 2001 From: Hugo Aguirre Parra Date: Thu, 1 May 2025 01:43:24 +0000 Subject: [PATCH 3/4] disable thinking, ignore thoughts --- go/plugins/googlegenai/gemini.go | 10 +++----- go/plugins/googlegenai/googleai_live_test.go | 26 +++++++++++++++++--- 2 files changed, 26 insertions(+), 10 deletions(-) diff --git a/go/plugins/googlegenai/gemini.go b/go/plugins/googlegenai/gemini.go index 4e0f38acb..c979ebae5 100644 --- a/go/plugins/googlegenai/gemini.go +++ b/go/plugins/googlegenai/gemini.go @@ -399,7 +399,6 @@ func generate( if err != nil { return nil, err } - fmt.Printf("thinking tokens: %d", resp.UsageMetadata.ThoughtsTokenCount) r := translateResponse(resp) r.Request = input if cache != nil { @@ -540,7 +539,6 @@ func toGeminiRequest(input *ai.ModelRequest, cache *genai.CachedContent) (*genai } if c.ThinkingConfig != nil { - fmt.Printf("setting thinking config!\n\n") gcc.ThinkingConfig = &genai.ThinkingConfig{ IncludeThoughts: c.ThinkingConfig.IncludeThoughts, ThinkingBudget: &c.ThinkingConfig.ThinkingBudget, @@ -786,10 +784,11 @@ func translateCandidate(cand *genai.Candidate) *ai.ModelResponse { if part.Text != "" { partFound++ - p = ai.NewTextPart(part.Text) if part.Thought { - fmt.Printf("this is a text part, and is a thought!\n\n part: %q", part.Text) + // TODO: Include a `reasoning` part. Not available in the SDK yet. + continue } + p = ai.NewTextPart(part.Text) } if part.InlineData != nil { partFound++ @@ -820,9 +819,6 @@ func translateCandidate(cand *genai.Candidate) *ai.ModelResponse { part.ExecutableCode.Code, ) } - if part.Thought { - fmt.Printf("THOUGHT found!!\n\n") - } if partFound > 1 { panic(fmt.Sprintf("expected only 1 content part in response, got %d, part: %#v", partFound, part)) } diff --git a/go/plugins/googlegenai/googleai_live_test.go b/go/plugins/googlegenai/googleai_live_test.go index 6b346e8cc..5c51cd164 100644 --- a/go/plugins/googlegenai/googleai_live_test.go +++ b/go/plugins/googlegenai/googleai_live_test.go @@ -416,9 +416,29 @@ func TestGoogleAILive(t *testing.T) { if resp == nil { t.Fatal("nil response obtanied") } - // since Thinking was enabled, the response should have thinking parts - fmt.Print(resp.Text()) - t.Fatal("remove me") + // TODO: add usageMetadata validation when SDK provides int + // see https://github.com/googleapis/go-genai/issues/282 + }) + 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") + } + // TODO: add usageMetadata validation when SDK provides int + // see https://github.com/googleapis/go-genai/issues/282 }) } From 9024b4cd97df45d51bafb548ccddc21617b62d39 Mon Sep 17 00:00:00 2001 From: Hugo Aguirre Parra Date: Thu, 8 May 2025 00:20:41 +0000 Subject: [PATCH 4/4] add test cases and minor input validations --- go/plugins/googlegenai/gemini.go | 15 +++++++++++--- go/plugins/googlegenai/gemini_test.go | 21 ++++++++++++++++++++ go/plugins/googlegenai/googleai_live_test.go | 12 ++++++----- 3 files changed, 40 insertions(+), 8 deletions(-) diff --git a/go/plugins/googlegenai/gemini.go b/go/plugins/googlegenai/gemini.go index 98e81c9c5..b3d3daa71 100644 --- a/go/plugins/googlegenai/gemini.go +++ b/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{ @@ -400,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) @@ -539,6 +545,9 @@ 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, diff --git a/go/plugins/googlegenai/gemini_test.go b/go/plugins/googlegenai/gemini_test.go index ce0bc6464..b00b9e8bc 100644 --- a/go/plugins/googlegenai/gemini_test.go +++ b/go/plugins/googlegenai/gemini_test.go @@ -151,6 +151,27 @@ func TestConvertRequest(t *testing.T) { 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} gt, err := toGeminiTools(tools) diff --git a/go/plugins/googlegenai/googleai_live_test.go b/go/plugins/googlegenai/googleai_live_test.go index 828f8ec9c..03ab13165 100644 --- a/go/plugins/googlegenai/googleai_live_test.go +++ b/go/plugins/googlegenai/googleai_live_test.go @@ -408,7 +408,7 @@ func TestGoogleAILive(t *testing.T) { Temperature: 0.4, ThinkingConfig: &googlegenai.ThinkingConfig{ IncludeThoughts: true, - ThinkingBudget: 1024, + ThinkingBudget: 100, }, }), ai.WithModel(m), @@ -419,8 +419,9 @@ func TestGoogleAILive(t *testing.T) { if resp == nil { t.Fatal("nil response obtanied") } - // TODO: add usageMetadata validation when SDK provides int - // see https://github.com/googleapis/go-genai/issues/282 + 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") @@ -440,8 +441,9 @@ func TestGoogleAILive(t *testing.T) { if resp == nil { t.Fatal("nil response obtanied") } - // TODO: add usageMetadata validation when SDK provides int - // see https://github.com/googleapis/go-genai/issues/282 + if resp.Usage.ThoughtsTokens > 0 { + t.Fatal("thoughts tokens should be zero") + } }) }