支持流式输出思考过程与思考模式配置

2026-08-14 16:13:45 +08:00
parent f7023f7958
commit 068c5fccf1
3 changed files with 108 additions and 13 deletions
@@ -2,10 +2,12 @@ package bot

 import (
 	"context"
+	"encoding/json"
 	"fmt"

 	"github.com/openai/openai-go"
 	"github.com/openai/openai-go/option"
+	"github.com/openai/openai-go/shared"
 	"myaibot/internal/config"
 )

@@ -61,7 +63,27 @@ func (b *Bot) SwitchModel(id string) error {
 	return nil
 }

-func (b *Bot) Chat(ctx context.Context, userMsg string) (string, error) {
+func (b *Bot) SetThinking(v string) error {
+	if v != "enabled" && v != "disabled" {
+		return fmt.Errorf("无效值: %s（可选 enabled/disabled）", v)
+	}
+	b.provider.Thinking = v
+	return nil
+}
+
+func (b *Bot) SetEffort(v string) error {
+	if v != "low" && v != "high" && v != "max" {
+		return fmt.Errorf("无效值: %s（可选 low/high/max）", v)
+	}
+	b.provider.ReasoningEffort = v
+	return nil
+}
+
+func (b *Bot) ThinkingConfig() (string, string) {
+	return b.provider.Thinking, b.provider.ReasoningEffort
+}
+
+func (b *Bot) Chat(ctx context.Context, userMsg string, onReasoning, onContent func(string)) (string, error) {
 	if b.provider.APIKey == "" {
 		return "", fmt.Errorf("供应商 %s 未配置 api_key，请编辑 data/config.yaml", b.provider.Name)
 	}
@@ -70,14 +92,32 @@ func (b *Bot) Chat(ctx context.Context, userMsg string) (string, error) {
 	history = append(history, b.history...)
 	history = append(history, openai.UserMessage(userMsg))

-	stream := b.client().Chat.Completions.NewStreaming(ctx, openai.ChatCompletionNewParams{
+	params := openai.ChatCompletionNewParams{
 		Model:    b.model,
 		Messages: history,
-	})
+	}
+	if p := b.provider; p.ReasoningEffort != "" && p.Thinking != "disabled" {
+		params.ReasoningEffort = shared.ReasoningEffort(p.ReasoningEffort)
+	}
+	if b.provider.Thinking != "" {
+		params.SetExtraFields(map[string]any{
+			"thinking": map[string]string{"type": b.provider.Thinking},
+		})
+	}
+	stream := b.client().Chat.Completions.NewStreaming(ctx, params)
 	answer := ""
 	for stream.Next() {
-		for _, delta := range stream.Current().Choices {
-			answer += delta.Delta.Content
+		for _, choice := range stream.Current().Choices {
+			if rc, ok := choice.Delta.JSON.ExtraFields["reasoning_content"]; ok && rc.Valid() {
+				var s string
+				if json.Unmarshal([]byte(rc.Raw()), &s) == nil && s != "" {
+					onReasoning(s)
+				}
+			}
+			if choice.Delta.Content != "" {
+				answer += choice.Delta.Content
+				onContent(choice.Delta.Content)
+			}
 		}
 	}
 	if err := stream.Err(); err != nil {
@@ -16,10 +16,12 @@ const (
 )

 type Provider struct {
-	Name    string   `yaml:"name"`
-	APIKey  string   `yaml:"api_key"`
-	BaseURL string   `yaml:"base_url"`
-	Models  []string `yaml:"models"`
+	Name            string   `yaml:"name"`
+	APIKey          string   `yaml:"api_key"`
+	BaseURL         string   `yaml:"base_url"`
+	Models          []string `yaml:"models"`
+	Thinking        string   `yaml:"thinking"`
+	ReasoningEffort string   `yaml:"reasoning_effort"`
 }

 type Config struct {
@@ -146,6 +148,12 @@ func validate(c *Config) error {
 				return fmt.Errorf("供应商 %s 包含空模型名", p.Name)
 			}
 		}
+		if p.Thinking != "" && !contains([]string{"enabled", "disabled"}, p.Thinking) {
+			return fmt.Errorf("供应商 %s 的 thinking 无效: %q（可选 enabled/disabled）", p.Name, p.Thinking)
+		}
+		if p.ReasoningEffort != "" && !contains([]string{"low", "high", "max"}, p.ReasoningEffort) {
+			return fmt.Errorf("供应商 %s 的 reasoning_effort 无效: %q（可选 low/high/max）", p.Name, p.ReasoningEffort)
+		}
 	}
 	if _, ok := names[c.DefaultProvider]; !ok {
 		return fmt.Errorf("default_provider %q 不存在", c.DefaultProvider)
@@ -38,12 +38,29 @@ func main() {
 			}
 			continue
 		}
-		answer, err := b.Chat(context.Background(), input)
+		fmt.Printf("%s: ", cfg.BotName)
+		thinkStyle, resetStyle := false, false
+		_, err = b.Chat(context.Background(), input,
+			func(text string) {
+				if !thinkStyle {
+					fmt.Print("\x1b[3;90m🧠 ")
+					thinkStyle, resetStyle = true, true
+				}
+				fmt.Print(text)
+			},
+			func(text string) {
+				if resetStyle {
+					fmt.Print("\x1b[0m")
+					thinkStyle, resetStyle = false, false
+				}
+				fmt.Print(text)
+			},
+		)
+		fmt.Println()
 		if err != nil {
 			fmt.Printf("⚠️  %v\n", err)
 			continue
 		}
-		fmt.Printf("%s: %s\n", cfg.BotName, answer)
 	}
 	if err := scanner.Err(); err != nil {
 		log.Fatal(err)
@@ -61,7 +78,9 @@ func handleCommand(b *bot.Bot, input string) bool {
 		fmt.Println("命令列表:")
 		fmt.Println("  /models            列出所有供应商和模型")
 		fmt.Println("  /use <模型>        切换模型，如 /use deepseek-chat 或 /use deepseek/deepseek-chat")
-		fmt.Println("  /info              显示当前供应商和模型")
+		fmt.Println("  /think <on|off>    开启或关闭当前供应商的思考模式")
+		fmt.Println("  /effort <low|high|max>  设置思考强度")
+		fmt.Println("  /info              显示当前供应商、模型和思考配置")
 		fmt.Println("  /exit              退出")
 	case "/models":
 		for _, m := range b.Models() {
@@ -78,9 +97,37 @@ func handleCommand(b *bot.Bot, input string) bool {
 		}
 		provider, model := b.Current()
 		fmt.Printf("已切换到 %s/%s (对话历史已保留)\n", provider, model)
+	case "/think":
+		if len(args) == 0 {
+			fmt.Println("用法: /think <on|off>")
+			return true
+		}
+		v := map[string]string{"on": "enabled", "off": "disabled"}[args[0]]
+		if err := b.SetThinking(v); err != nil {
+			fmt.Printf("⚠️  %v\n", err)
+			return true
+		}
+		fmt.Printf("思考模式已%s\n", map[string]string{"enabled": "开启", "disabled": "关闭"}[v])
+	case "/effort":
+		if len(args) == 0 {
+			fmt.Println("用法: /effort <low|high|max>")
+			return true
+		}
+		if err := b.SetEffort(args[0]); err != nil {
+			fmt.Printf("⚠️  %v\n", err)
+			return true
+		}
+		fmt.Printf("思考强度已设置为 %s\n", args[0])
 	case "/info":
 		provider, model := b.Current()
-		fmt.Printf("供应商: %s, 模型: %s\n", provider, model)
+		thinking, effort := b.ThinkingConfig()
+		if thinking == "" {
+			thinking = "enabled(默认)"
+		}
+		if effort == "" {
+			effort = "high(默认)"
+		}
+		fmt.Printf("供应商: %s, 模型: %s, 思考模式: %s, 思考强度: %s\n", provider, model, thinking, effort)
 	default:
 		fmt.Printf("未知命令: %s，输入 /help 查看命令列表\n", cmd)
 	}