From 1c9e98ddb0b9b444d4cd3957f4dfdfe79f13bacb Mon Sep 17 00:00:00 2001 From: kevin Date: Fri, 14 Aug 2026 16:55:48 +0800 Subject: [PATCH] =?UTF-8?q?=E6=94=AF=E6=8C=81=E5=B7=A5=E5=85=B7=E8=B0=83?= =?UTF-8?q?=E7=94=A8=EF=BC=9A=E4=B8=89=E4=B8=AAAI=E8=A7=92=E8=89=B2?= =?UTF-8?q?=E9=85=8D=E7=BD=AE=E3=80=81=E5=B7=A5=E5=85=B7=E6=B3=A8=E5=86=8C?= =?UTF-8?q?=E8=A1=A8=E4=B8=8E=E6=B5=81=E5=BC=8F=E5=B7=A5=E5=85=B7=E8=BD=AE?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- go.mod | 3 +- go.sum | 2 + internal/bot/bot.go | 202 ++++++++++++++++++++++++++---- internal/cli/cli.go | 2 + internal/config/config.go | 12 ++ internal/tools/calculator_tool.go | 61 +++++++++ internal/tools/random_tool.go | 42 +++++++ internal/tools/time_tool.go | 38 ++++++ internal/tools/tools.go | 56 +++++++++ internal/tools/tools_test.go | 87 +++++++++++++ main.go | 14 +++ 11 files changed, 491 insertions(+), 28 deletions(-) create mode 100644 internal/tools/calculator_tool.go create mode 100644 internal/tools/random_tool.go create mode 100644 internal/tools/time_tool.go create mode 100644 internal/tools/tools.go create mode 100644 internal/tools/tools_test.go diff --git a/go.mod b/go.mod index 0578e76..abac104 100644 --- a/go.mod +++ b/go.mod @@ -3,14 +3,15 @@ module myaibot go 1.26.4 require ( + github.com/expr-lang/expr v1.17.8 github.com/openai/openai-go v1.12.0 github.com/peterh/liner v1.2.2 + github.com/tidwall/gjson v1.14.4 gopkg.in/yaml.v3 v3.0.1 ) require ( github.com/mattn/go-runewidth v0.0.3 // indirect - github.com/tidwall/gjson v1.14.4 // indirect github.com/tidwall/match v1.1.1 // indirect github.com/tidwall/pretty v1.2.1 // indirect github.com/tidwall/sjson v1.2.5 // indirect diff --git a/go.sum b/go.sum index 1db1bdf..11d6d28 100644 --- a/go.sum +++ b/go.sum @@ -1,3 +1,5 @@ +github.com/expr-lang/expr v1.17.8 h1:W1loDTT+0PQf5YteHSTpju2qfUfNoBt4yw9+wOEU9VM= +github.com/expr-lang/expr v1.17.8/go.mod h1:8/vRC7+7HBzESEqt5kKpYXxrxkr31SaO8r40VO/1IT4= github.com/mattn/go-runewidth v0.0.3 h1:a+kO+98RDGEfo6asOGMmpodZq4FNtnGP54yps8BzLR4= github.com/mattn/go-runewidth v0.0.3/go.mod h1:LwmH8dsx7+W8Uxz3IHJYH5QSwggIsqBzpuz5H//U1FU= github.com/openai/openai-go v1.12.0 h1:NBQCnXzqOTv5wsgNC36PrFEiskGfO5wccfCWDo9S1U0= diff --git a/internal/bot/bot.go b/internal/bot/bot.go index 8bf8c40..8b6a9aa 100644 --- a/internal/bot/bot.go +++ b/internal/bot/bot.go @@ -8,19 +8,32 @@ import ( "github.com/openai/openai-go" "github.com/openai/openai-go/option" + "github.com/openai/openai-go/packages/param" "github.com/openai/openai-go/shared" + "github.com/openai/openai-go/shared/constant" + "github.com/tidwall/gjson" "myaibot/internal/config" + "myaibot/internal/tools" ) -const maxHistory = 20 +const ( + maxHistory = 20 + maxToolRounds = 5 +) + +var toolRegistry = tools.NewRegistry(tools.TimeTool{}, tools.CalculatorTool{}, tools.RandomTool{}) type Bot struct { - clients map[string]*openai.Client - cfg *config.Config - provider *config.Provider - model string - history []openai.ChatCompletionMessageParamUnion - systemPrompt string + clients map[string]*openai.Client + cfg *config.Config + provider *config.Provider + model string + toolProvider *config.Provider + toolModel string + visionProvider *config.Provider + visionModel string + history []openai.ChatCompletionMessageParamUnion + systemPrompt string } func New(cfg *config.Config) *Bot { @@ -31,21 +44,37 @@ func New(cfg *config.Config) *Bot { } b.provider = config.FindProvider(cfg.DefaultProvider) b.model = cfg.DefaultModel + b.toolProvider, b.toolModel = b.provider, b.model + if cfg.ToolModel != "" { + if p, m, err := config.ResolveModel(cfg.ToolModel); err == nil { + b.toolProvider, b.toolModel = p, m + } + } + b.visionProvider, b.visionModel = b.provider, b.model + if cfg.VisionModel != "" { + if p, m, err := config.ResolveModel(cfg.VisionModel); err == nil { + b.visionProvider, b.visionModel = p, m + } + } return b } -func (b *Bot) client() *openai.Client { - if c, ok := b.clients[b.provider.Name]; ok { +func (b *Bot) clientFor(p *config.Provider) *openai.Client { + if c, ok := b.clients[p.Name]; ok { return c } c := openai.NewClient( - option.WithAPIKey(b.provider.APIKey), - option.WithBaseURL(b.provider.BaseURL), + option.WithAPIKey(p.APIKey), + option.WithBaseURL(p.BaseURL), ) - b.clients[b.provider.Name] = &c + b.clients[p.Name] = &c return &c } +func (b *Bot) client() *openai.Client { + return b.clientFor(b.provider) +} + func (b *Bot) Models() []string { return config.AllModels() } @@ -84,6 +113,11 @@ func (b *Bot) ThinkingConfig() (string, string) { return b.provider.Thinking, b.provider.ReasoningEffort } +func (b *Bot) CurrentRoles() (tool, vision string) { + return b.toolProvider.Name + "/" + b.toolModel, + b.visionProvider.Name + "/" + b.visionModel +} + func (b *Bot) ContextDump() string { var sb strings.Builder sb.WriteString("[系统] " + b.systemPrompt + "\n") @@ -105,36 +139,33 @@ func (b *Bot) ContextDump() string { return sb.String() } -func (b *Bot) Chat(ctx context.Context, userMsg string, onReasoning, onContent func(string)) (string, error) { +func (b *Bot) Chat(ctx context.Context, userMsg string, onReasoning, onContent func(string), onTool func(string, string), onToolReasoning func(string)) (string, error) { if b.provider.APIKey == "" { return "", fmt.Errorf("供应商 %s 未配置 api_key，请编辑 data/config.yaml", b.provider.Name) } - history := make([]openai.ChatCompletionMessageParamUnion, 0, len(b.history)+2) + toolMsgs, usedTool, err := b.toolRound(ctx, userMsg, onToolReasoning, onTool) + if err != nil { + return "", err + } + history := make([]openai.ChatCompletionMessageParamUnion, 0, len(b.history)+len(toolMsgs)+2) history = append(history, openai.SystemMessage(b.systemPrompt)) history = append(history, b.history...) history = append(history, openai.UserMessage(userMsg)) + if usedTool { + history = append(history, toolMsgs...) + } 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}, - }) - } + b.applyThinkingParams(¶ms, b.provider) stream := b.client().Chat.Completions.NewStreaming(ctx, params) answer := "" for stream.Next() { 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 s := extraReasoning(choice.Delta.RawJSON()); s != "" { + onReasoning(s) } if choice.Delta.Content != "" { answer += choice.Delta.Content @@ -151,3 +182,120 @@ func (b *Bot) Chat(ctx context.Context, userMsg string, onReasoning, onContent f } return answer, nil } + +func (b *Bot) toolRound(ctx context.Context, userMsg string, onReasoning func(string), onTool func(string, string)) ([]openai.ChatCompletionMessageParamUnion, bool, error) { + if b.toolProvider.APIKey == "" { + return nil, false, fmt.Errorf("工具调用供应商 %s 未配置 api_key，请编辑 data/config.yaml", b.toolProvider.Name) + } + history := make([]openai.ChatCompletionMessageParamUnion, 0, len(b.history)+2) + history = append(history, openai.SystemMessage(b.systemPrompt)) + history = append(history, b.history...) + history = append(history, openai.UserMessage(userMsg)) + + var toolMsgs []openai.ChatCompletionMessageParamUnion + for range maxToolRounds { + params := openai.ChatCompletionNewParams{ + Model: b.toolModel, + Messages: history, + Tools: toolRegistry.ParamList(), + } + b.applyThinkingParams(¶ms, b.toolProvider) + stream := b.clientFor(b.toolProvider).Chat.Completions.NewStreaming(ctx, params) + + var ( + reasoning strings.Builder + calls []pendingToolCall + ) + for stream.Next() { + for _, choice := range stream.Current().Choices { + if s := extraReasoning(choice.Delta.RawJSON()); s != "" { + reasoning.WriteString(s) + onReasoning(s) + } + for _, tc := range choice.Delta.ToolCalls { + for len(calls) <= int(tc.Index) { + calls = append(calls, pendingToolCall{}) + } + c := &calls[tc.Index] + if tc.ID != "" { + c.id = tc.ID + } + if tc.Function.Name != "" { + c.name = tc.Function.Name + } + c.arguments.WriteString(tc.Function.Arguments) + } + } + } + if err := stream.Err(); err != nil { + return nil, false, fmt.Errorf("工具调用 AI 请求失败: %w", err) + } + if len(calls) == 0 { + break + } + asst := openai.ChatCompletionAssistantMessageParam{ + Content: openai.ChatCompletionAssistantMessageParamContentUnion{ + OfString: param.NewOpt(""), + }, + ToolCalls: pendingToParams(calls), + } + if s := reasoning.String(); s != "" { + asst.SetExtraFields(map[string]any{"reasoning_content": s}) + } + asstMsg := openai.ChatCompletionMessageParamUnion{OfAssistant: &asst} + history = append(history, asstMsg) + toolMsgs = append(toolMsgs, asstMsg) + for _, tc := range asst.ToolCalls { + if onTool != nil { + onTool(tc.Function.Name, tc.Function.Arguments) + } + result, err := toolRegistry.Execute(tc.Function.Name, json.RawMessage(tc.Function.Arguments)) + if err != nil { + result = "工具执行失败: " + err.Error() + } + toolMsg := openai.ToolMessage(result, tc.ID) + history = append(history, toolMsg) + toolMsgs = append(toolMsgs, toolMsg) + } + } + return toolMsgs, len(toolMsgs) > 0, nil +} + +func (b *Bot) applyThinkingParams(params *openai.ChatCompletionNewParams, p *config.Provider) { + if p.ReasoningEffort != "" && p.Thinking != "disabled" { + params.ReasoningEffort = shared.ReasoningEffort(p.ReasoningEffort) + } + if p.Thinking != "" { + params.SetExtraFields(map[string]any{ + "thinking": map[string]string{"type": p.Thinking}, + }) + } +} + +func extraReasoning(raw string) string { + if raw == "" { + return "" + } + return gjson.Get(raw, "reasoning_content").String() +} + +type pendingToolCall struct { + id string + name string + arguments strings.Builder +} + +func pendingToParams(calls []pendingToolCall) []openai.ChatCompletionMessageToolCallParam { + out := make([]openai.ChatCompletionMessageToolCallParam, 0, len(calls)) + for _, c := range calls { + out = append(out, openai.ChatCompletionMessageToolCallParam{ + ID: c.id, + Function: openai.ChatCompletionMessageToolCallFunctionParam{ + Name: c.name, + Arguments: c.arguments.String(), + }, + Type: constant.Function("function"), + }) + } + return out +} diff --git a/internal/cli/cli.go b/internal/cli/cli.go index cb89a6d..35fc160 100644 --- a/internal/cli/cli.go +++ b/internal/cli/cli.go @@ -72,6 +72,7 @@ func (h *Handler) Handle(input string) bool { case "/info": provider, model := h.bot.Current() thinking, effort := h.bot.ThinkingConfig() + tool, vision := h.bot.CurrentRoles() if thinking == "" { thinking = "enabled(默认)" } @@ -79,6 +80,7 @@ func (h *Handler) Handle(input string) bool { effort = "high(默认)" } fmt.Printf("供应商: %s, 模型: %s, 思考模式: %s, 思考强度: %s\n", provider, model, thinking, effort) + fmt.Printf("工具调用AI: %s\n图片识别AI: %s\n", tool, vision) default: fmt.Printf("未知命令: %s，输入 /help 查看命令列表\n", cmd) } diff --git a/internal/config/config.go b/internal/config/config.go index 528d214..d5a97dd 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -32,6 +32,8 @@ type Config struct { Providers []Provider `yaml:"providers"` DefaultProvider string `yaml:"default_provider"` DefaultModel string `yaml:"default_model"` + ToolModel string `yaml:"tool_model"` + VisionModel string `yaml:"vision_model"` } type legacyConfig struct { @@ -161,6 +163,16 @@ func validate(c *Config) error { if _, _, err := ResolveModel(c.DefaultModel); err != nil { return fmt.Errorf("default_model 无效: %w", err) } + if c.ToolModel != "" { + if _, _, err := ResolveModel(c.ToolModel); err != nil { + return fmt.Errorf("tool_model 无效: %w", err) + } + } + if c.VisionModel != "" { + if _, _, err := ResolveModel(c.VisionModel); err != nil { + return fmt.Errorf("vision_model 无效: %w", err) + } + } return nil } diff --git a/internal/tools/calculator_tool.go b/internal/tools/calculator_tool.go new file mode 100644 index 0000000..ac48603 --- /dev/null +++ b/internal/tools/calculator_tool.go @@ -0,0 +1,61 @@ +package tools + +import ( + "encoding/json" + "fmt" + "math" + "strconv" + + "github.com/expr-lang/expr" +) + +type CalculatorTool struct{} + +func (CalculatorTool) Name() string { return "calculate" } +func (CalculatorTool) Description() string { + return "计算数学表达式，如 \"(12 + 3) * 4\"、\"2^10\"、\"sqrt(9)\"" +} +func (CalculatorTool) Parameters() map[string]any { + return map[string]any{ + "type": "object", + "properties": map[string]any{ + "expression": map[string]any{"type": "string", "description": "要计算的数学表达式"}, + }, + "required": []string{"expression"}, + } +} + +func (CalculatorTool) Execute(args json.RawMessage) (string, error) { + var p struct { + Expression string `json:"expression"` + } + if err := json.Unmarshal(args, &p); err != nil { + return "", err + } + if p.Expression == "" { + return "", fmt.Errorf("expression 不能为空") + } + out, err := expr.Eval(p.Expression, nil) + if err != nil { + return "", err + } + switch v := out.(type) { + case float64: + return strconv.FormatFloat(round(v), 'f', -1, 64), nil + case int: + return strconv.Itoa(v), nil + case bool: + return strconv.FormatBool(v), nil + case string: + return v, nil + default: + return fmt.Sprintf("%v", v), nil + } +} + +func round(v float64) float64 { + if math.IsInf(v, 0) || math.IsNaN(v) { + return v + } + return math.Round(v*1e8) / 1e8 +} diff --git a/internal/tools/random_tool.go b/internal/tools/random_tool.go new file mode 100644 index 0000000..7c74cc3 --- /dev/null +++ b/internal/tools/random_tool.go @@ -0,0 +1,42 @@ +package tools + +import ( + "encoding/json" + "fmt" + "math/rand/v2" +) + +type RandomTool struct{} + +func (RandomTool) Name() string { return "random_number" } +func (RandomTool) Description() string { + return "生成指定范围内的随机整数，默认 0 到 100（含端点）" +} +func (RandomTool) Parameters() map[string]any { + return map[string]any{ + "type": "object", + "properties": map[string]any{ + "min": map[string]any{"type": "integer", "description": "最小值，默认 0"}, + "max": map[string]any{"type": "integer", "description": "最大值，默认 100"}, + }, + } +} + +func (RandomTool) Execute(args json.RawMessage) (string, error) { + var p struct { + Min *int `json:"min"` + Max *int `json:"max"` + } + _ = json.Unmarshal(args, &p) + min, max := 0, 100 + if p.Min != nil { + min = *p.Min + } + if p.Max != nil { + max = *p.Max + } + if max < min { + return "", fmt.Errorf("max (%d) 不能小于 min (%d)", max, min) + } + return fmt.Sprintf("%d", rand.IntN(max-min+1)+min), nil +} diff --git a/internal/tools/time_tool.go b/internal/tools/time_tool.go new file mode 100644 index 0000000..3c259c9 --- /dev/null +++ b/internal/tools/time_tool.go @@ -0,0 +1,38 @@ +package tools + +import ( + "encoding/json" + "time" +) + +type TimeTool struct{} + +func (TimeTool) Name() string { return "get_current_time" } +func (TimeTool) Description() string { + return "获取当前日期和时间，可选指定时区（如 Asia/Shanghai，默认本地时区）" +} +func (TimeTool) Parameters() map[string]any { + return map[string]any{ + "type": "object", + "properties": map[string]any{ + "timezone": map[string]any{"type": "string", "description": "IANA 时区名，如 Asia/Shanghai"}, + }, + } +} + +func (TimeTool) Execute(args json.RawMessage) (string, error) { + var p struct { + Timezone string `json:"timezone"` + } + _ = json.Unmarshal(args, &p) + loc := time.Local + if p.Timezone != "" { + l, err := time.LoadLocation(p.Timezone) + if err != nil { + return "", err + } + loc = l + } + now := time.Now().In(loc) + return now.Format("2006-01-02 15:04:05 Monday MST"), nil +} diff --git a/internal/tools/tools.go b/internal/tools/tools.go new file mode 100644 index 0000000..a48180c --- /dev/null +++ b/internal/tools/tools.go @@ -0,0 +1,56 @@ +package tools + +import ( + "encoding/json" + "fmt" + + "github.com/openai/openai-go" + "github.com/openai/openai-go/packages/param" + "github.com/openai/openai-go/shared" +) + +type Tool interface { + Name() string + Description() string + Parameters() map[string]any + Execute(args json.RawMessage) (string, error) +} + +type Registry struct { + tools map[string]Tool +} + +func NewRegistry(t ...Tool) *Registry { + r := &Registry{tools: make(map[string]Tool)} + for _, tool := range t { + r.tools[tool.Name()] = tool + } + return r +} + +func (r *Registry) Get(name string) (Tool, bool) { + t, ok := r.tools[name] + return t, ok +} + +func (r *Registry) ParamList() []openai.ChatCompletionToolParam { + out := make([]openai.ChatCompletionToolParam, 0, len(r.tools)) + for _, t := range r.tools { + out = append(out, openai.ChatCompletionToolParam{ + Function: shared.FunctionDefinitionParam{ + Name: t.Name(), + Description: param.NewOpt(t.Description()), + Parameters: shared.FunctionParameters(t.Parameters()), + }, + }) + } + return out +} + +func (r *Registry) Execute(name string, args json.RawMessage) (string, error) { + t, ok := r.tools[name] + if !ok { + return "", fmt.Errorf("未知工具: %s", name) + } + return t.Execute(args) +} diff --git a/internal/tools/tools_test.go b/internal/tools/tools_test.go new file mode 100644 index 0000000..906e584 --- /dev/null +++ b/internal/tools/tools_test.go @@ -0,0 +1,87 @@ +package tools + +import ( + "encoding/json" + "strings" + "testing" +) + +func TestCalculator(t *testing.T) { + cases := []struct { + name string + expr string + want string + }{ + {"四则运算", "(12 + 3) * 4", "60"}, + {"幂运算", "2^10", "1024"}, + {"浮点", "7 / 2", "3.5"}, + {"小数精度", "1 / 3", "0.33333333"}, + {"布尔", "2 > 1", "true"}, + } + for _, c := range cases { + t.Run(c.name, func(t *testing.T) { + args, _ := json.Marshal(map[string]string{"expression": c.expr}) + got, err := CalculatorTool{}.Execute(args) + if err != nil { + t.Fatalf("Execute 出错: %v", err) + } + if got != c.want { + t.Errorf("Execute(%q) = %q, want %q", c.expr, got, c.want) + } + }) + } +} + +func TestCalculatorInvalid(t *testing.T) { + args, _ := json.Marshal(map[string]string{"expression": "1 +"}) + if _, err := (CalculatorTool{}).Execute(args); err == nil { + t.Error("非法表达式应返回错误") + } +} + +func TestRandom(t *testing.T) { + args, _ := json.Marshal(map[string]any{"min": 1, "max": 10}) + for i := 0; i < 100; i++ { + out, err := RandomTool{}.Execute(args) + if err != nil { + t.Fatalf("Execute 出错: %v", err) + } + var v int + if err := json.Unmarshal([]byte(out), &v); err != nil { + t.Fatalf("结果 %q 解析失败: %v", out, err) + } + if v < 1 || v > 10 { + t.Fatalf("结果 %q 不在 [1,10] 内", out) + } + } + bad, _ := json.Marshal(map[string]any{"min": 10, "max": 1}) + if _, err := (RandomTool{}).Execute(bad); err == nil { + t.Error("min>max 应返回错误") + } +} + +func TestTimeTool(t *testing.T) { + out, err := TimeTool{}.Execute(nil) + if err != nil { + t.Fatalf("Execute 出错: %v", err) + } + if !strings.Contains(out, "20") { + t.Errorf("时间输出异常: %q", out) + } +} + +func TestRegistry(t *testing.T) { + r := NewRegistry(TimeTool{}, CalculatorTool{}, RandomTool{}) + if _, ok := r.Get("get_current_time"); !ok { + t.Error("get_current_time 未注册") + } + if _, ok := r.Get("nonexistent"); ok { + t.Error("未知工具不应存在") + } + if len(r.ParamList()) != 3 { + t.Errorf("ParamList 数量 = %d, want 3", len(r.ParamList())) + } + if _, err := r.Execute("nonexistent", nil); err == nil { + t.Error("执行未知工具应返回错误") + } +} diff --git a/main.go b/main.go index 11ceccf..fbe4616 100644 --- a/main.go +++ b/main.go @@ -67,6 +67,20 @@ func main() { } fmt.Print(text) }, + func(name, args string) { + if resetStyle { + fmt.Print("\x1b[0m") + thinkStyle, resetStyle = false, false + } + fmt.Printf("🔧 调用工具: %s %s\n", name, args) + }, + func(text string) { + if !thinkStyle { + fmt.Print("\x1b[3;90m🔧 ") + thinkStyle, resetStyle = true, true + } + fmt.Print(text) + }, ) fmt.Println() if err != nil {