feat: enhance llm.Config with ThinkingConfig and add corresponding test

Author: lemon-mint

llm/llm_configs.go

@@ -1,14 +1,20 @@
 package llm
 
 type Config struct {
-	Temperature     *float32 `json:"temperature"`
-	TopP            *float32 `json:"top_p"`
-	TopK            *int     `json:"top_k"`
-	MaxOutputTokens *int     `json:"max_output_tokens"`
-	StopSequences   []string `json:"stop_sequences"`
+	Temperature     *float32        `json:"temperature,omitempty"`
+	TopP            *float32        `json:"top_p,omitempty"`
+	TopK            *int            `json:"top_k,omitempty"`
+	MaxOutputTokens *int            `json:"max_output_tokens,omitempty"`
+	StopSequences   []string        `json:"stop_sequences,omitempty"`
+	ThinkingConfig  *ThinkingConfig `json:"thinking_config,omitempty"`
 
-	SystemInstruction     string         `json:"system_instruction"`
-	SafetyFilterThreshold BlockThreshold `json:"filter_threshold"`
+	SystemInstruction     string         `json:"system_instruction,omitempty"`
+	SafetyFilterThreshold BlockThreshold `json:"filter_threshold,omitempty"`
+}
+
+type ThinkingConfig struct {
+	IncludeThoughts *bool `json:"include_thoughts,omitempty"`
+	ThinkingBudget  *int  `json:"thinking_budget,omitempty"`
 }
 
 type BlockThreshold uint16

pconf/utils.go

@@ -0,0 +1,5 @@
+package pconf
+
+func Ptrify[T any](v T) *T {
+	return &v
+}

provider/aistudio/aistudio_model.go

@@ -273,6 +273,16 @@ func (g *generativeLanguageModel) GenerateStream(ctx context.Context, chat *llm.
 		config.SystemInstruction = &genai.Content{Parts: []*genai.Part{{Text: g.config.SystemInstruction + chat.SystemInstruction}}}
 	}
 
+	if g.config.ThinkingConfig != nil {
+		config.ThinkingConfig = &genai.ThinkingConfig{}
+		if g.config.ThinkingConfig.IncludeThoughts != nil {
+			config.ThinkingConfig.IncludeThoughts = *g.config.ThinkingConfig.IncludeThoughts
+		}
+		if g.config.ThinkingConfig.ThinkingBudget != nil {
+			config.ThinkingConfig.ThinkingBudget = ptrify(int32(*g.config.ThinkingConfig.ThinkingBudget))
+		}
+	}
+
 	stream := make(chan llm.Segment, 128)
 	v := &llm.StreamContent{
 		Content: &llm.Content{},

provider/aistudio/aistudio_model_test.go

@@ -207,3 +207,39 @@ func TestAIStudioToolCall(t *testing.T) {
 		return
 	}
 }
+
+func TestAIStudioGenerateWithThinking(t *testing.T) {
+	client := getClient()
+	defer client.Close()
+
+	config := &llm.Config{
+		ThinkingConfig: &llm.ThinkingConfig{
+			ThinkingBudget: pconf.Ptrify(0),
+		},
+	}
+
+	model, err := client.NewLLM("gemini-2.5-flash-preview-04-17", config)
+	if err != nil {
+		panic(err)
+	}
+	defer model.Close()
+
+	output := model.GenerateStream(
+		context.Background(),
+		nil,
+		&llm.Content{
+			Role:  llm.RoleUser,
+			Parts: []llm.Segment{llm.Text("Do hamsters eat cats!")},
+		},
+	)
+
+	for segment := range output.Stream {
+		fmt.Print(segment)
+	}
+	fmt.Println()
+
+	if output.Err != nil {
+		t.Error(output.Err)
+		return
+	}
+}