kevin/ai-bot
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

支持工具调用：三个AI角色配置、工具注册表与流式工具轮

Browse Source
This commit is contained in:
kevin
2026-08-14 16:55:48 +08:00
parent ce49bb9c1e
commit 1c9e98ddb0
11 changed files with 491 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal/bot/bot.go
+175 -27
View File
@@ -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(&params, 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(&params, 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
}
internal/cli/cli.go
+2
View File
@@ -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)
}
internal/config/config.go
+12
View File
@@ -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
}
internal/tools/calculator_tool.go
+61
View File
@@ -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
}
internal/tools/random_tool.go
+42
View File
@@ -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
}
internal/tools/time_tool.go
+38
View File
@@ -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
}
internal/tools/tools.go
+56
View File
@@ -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)
}
internal/tools/tools_test.go
+87
View File
@@ -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("执行未知工具应返回错误")
}
}