kevin/aichat
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

动态加载工具配置

Browse Source 
Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
kevin
2026-06-17 13:13:06 +08:00
co-authored by Claude
parent 249227ef0a
commit 4e4c6b32fc
13 changed files with 704 additions and 188 deletions
Download Patch File Download Diff File Expand all files Collapse all files
agents/calculator/calculator.go
+90
View File
@@ -1,15 +1,21 @@
package calculator
import (
"context"
"encoding/json"
"errors"
"fmt"
"math"
"os"
"path/filepath"
"strconv"
"strings"
"unicode"
agents "aichat/agenttool"
"github.com/volcengine/volcengine-go-sdk/service/arkruntime/model"
"gopkg.in/yaml.v3"
)
const (
@@ -17,11 +23,95 @@ const (
ActivationPrompt = "执行简单、确定性的数学四则运算。当用户询问加减乘除、括号表达式、小数运算或需要准确计算表达式结果时，应直接调用此工具；不用于代数推导、方程求解、统计分析或复杂数学证明。"
)
type Config struct {
Enabled bool `yaml:"enabled" json:"enabled"`
ActivationPrompt string `yaml:"activation_prompt" json:"activation_prompt"`
}
type ToolArgs struct {
Expression string `json:"expression"`
Reason string `json:"reason"`
}
type LoadedTool struct {
cfg *Config
}
func NewLoadedTool(cfg *Config) *LoadedTool {
if cfg == nil {
defaultCfg := defaultConfig()
cfg = &defaultCfg
}
return &LoadedTool{cfg: cfg}
}
func init() {
agents.Register(agents.Descriptor{Name: ToolName, Load: func(path string, options agents.LoadOptions) (agents.LoadedTool, error) {
cfg, err := LoadConfig(path)
if err != nil {
return nil, err
}
return NewLoadedTool(cfg), nil
}})
}
func defaultConfig() Config {
return Config{Enabled: true, ActivationPrompt: ActivationPrompt}
}
func LoadConfig(path string) (*Config, error) {
if _, err := os.Stat(path); err != nil {
if !os.IsNotExist(err) {
return nil, fmt.Errorf("检查计算器工具配置失败: %w", err)
}
cfg := defaultConfig()
if err := os.MkdirAll(filepath.Dir(path), 0755); err != nil {
return nil, fmt.Errorf("创建计算器工具配置目录失败: %w", err)
}
data, err := yaml.Marshal(&cfg)
if err != nil {
return nil, fmt.Errorf("生成计算器工具配置失败: %w", err)
}
if err := os.WriteFile(path, data, 0644); err != nil {
return nil, fmt.Errorf("写入计算器工具配置失败: %w", err)
}
}
data, err := os.ReadFile(path)
if err != nil {
return nil, fmt.Errorf("读取计算器工具配置失败: %w", err)
}
var cfg Config
if err := yaml.Unmarshal(data, &cfg); err != nil {
return nil, fmt.Errorf("解析计算器工具配置失败: %w", err)
}
if strings.TrimSpace(cfg.ActivationPrompt) == "" {
cfg.ActivationPrompt = ActivationPrompt
}
return &cfg, nil
}
func (t *LoadedTool) Name() string { return ToolName }
func (t *LoadedTool) Enabled() bool { return t != nil && t.cfg != nil && t.cfg.Enabled }
func (t *LoadedTool) ToolDefinition(description string) *model.Tool {
if strings.TrimSpace(description) == "" && t != nil && t.cfg != nil {
description = t.cfg.ActivationPrompt
}
return ToolDefinition(description)
}
func (t *LoadedTool) Execute(ctx context.Context, args string, runtime agents.Runtime) (string, error) {
result, err := ExecuteTool(args)
if err == nil && runtime.Emit != nil {
runtime.Emit(agents.Frame{Type: "trace", Tool: ToolName, Stage: "calculate", Status: "success", Message: "四则运算完成"})
}
return result, err
}
func (t *LoadedTool) RawState() any { return nil }
func ToolDefinition(description string) *model.Tool {
description = strings.TrimSpace(description)
if description == "" {
agents/search/search.go
+80
View File
@@ -14,6 +14,8 @@ import (
"sync"
"time"
agents "aichat/agenttool"
"github.com/volcengine/volcengine-go-sdk/service/arkruntime/model"
"gopkg.in/yaml.v3"
)
@@ -82,6 +84,10 @@ type State struct {
activeName string
}
type LoadedTool struct {
state *State
}
type Result struct {
Title string `json:"title"`
URL string `json:"url"`
@@ -98,6 +104,80 @@ type ToolArgs struct {
Reason string `json:"reason"`
}
func init() {
agents.Register(agents.Descriptor{Name: ToolName, Load: func(path string, options agents.LoadOptions) (agents.LoadedTool, error) {
legacy, err := legacyProfilesFromOption(options.Value(agents.LegacySearchProfilesKey))
if err != nil {
return nil, err
}
cfg, err := LoadConfig(path, legacy)
if err != nil {
return nil, err
}
state, err := NewState(cfg)
if err != nil {
return nil, err
}
return &LoadedTool{state: state}, nil
}})
}
func legacyProfilesFromOption(value any) ([]ProfileConfig, error) {
switch profiles := value.(type) {
case nil:
return nil, nil
case []ProfileConfig:
return profiles, nil
case []map[string]any:
data, err := yaml.Marshal(profiles)
if err != nil {
return nil, fmt.Errorf("转换旧搜索配置失败: %w", err)
}
var parsed []ProfileConfig
if err := yaml.Unmarshal(data, &parsed); err != nil {
return nil, fmt.Errorf("解析旧搜索配置失败: %w", err)
}
return parsed, nil
default:
return nil, nil
}
}
func (t *LoadedTool) Name() string { return ToolName }
func (t *LoadedTool) Enabled() bool { return t != nil && t.state != nil && t.state.Enabled() }
func (t *LoadedTool) ToolDefinition(description string) *model.Tool {
if t == nil || t.state == nil {
return nil
}
return t.state.ToolDefinition(description)
}
func (t *LoadedTool) Execute(ctx context.Context, args string, runtime agents.Runtime) (string, error) {
if runtime.Emit != nil {
runtime.Emit(agents.Frame{Type: "trace", Tool: ToolName, Stage: "request", Status: "running", Message: "正在联网搜索"})
}
result, err := t.state.ExecuteTool(ctx, args)
if runtime.Emit != nil {
status := "success"
messageText := "联网搜索完成"
if err != nil {
status = "error"
messageText = "联网搜索失败"
}
runtime.Emit(agents.Frame{Type: "trace", Tool: ToolName, Stage: "results", Status: status, Message: messageText})
}
return result, err
}
func (t *LoadedTool) RawState() any {
if t == nil {
return nil
}
return t.state
}
type braveSearchResponse struct {
Web struct {
Results []Result `json:"results"`
agents/sql/sql_query.go
+66
View File
@@ -15,6 +15,8 @@ import (
"time"
"unicode/utf8"
agents "aichat/agenttool"
_ "github.com/go-sql-driver/mysql"
"github.com/volcengine/volcengine-go-sdk/service/arkruntime/model"
"gopkg.in/yaml.v3"
@@ -69,6 +71,10 @@ type State struct {
cacheAt time.Time
}
type LoadedTool struct {
state *State
}
type database struct {
cfg DatabaseConfig
db *sql.DB
@@ -90,6 +96,66 @@ type ToolArgs struct {
type SQLGenerator func(ctx context.Context, prompt string, maxTokens int) (string, error)
func init() {
agents.Register(agents.Descriptor{Name: ToolName, Load: func(path string, options agents.LoadOptions) (agents.LoadedTool, error) {
cfg, err := LoadConfig(path)
if err != nil {
return nil, err
}
state, err := NewState(cfg)
if err != nil {
return nil, err
}
return &LoadedTool{state: state}, nil
}})
}
func (t *LoadedTool) Name() string { return ToolName }
func (t *LoadedTool) Enabled() bool { return t != nil && t.state != nil && t.state.Enabled() }
func (t *LoadedTool) ToolDefinition(description string) *model.Tool {
if t == nil || t.state == nil {
return nil
}
return t.state.ToolDefinition(description)
}
func (t *LoadedTool) Execute(ctx context.Context, args string, runtime agents.Runtime) (string, error) {
if runtime.Emit != nil {
runtime.Emit(agents.Frame{Type: "trace", Tool: ToolName, Stage: "execute", Status: "running", Message: "正在查询数据库"})
}
generator := runtime.CompleteText
if generator == nil {
return "", errors.New("SQL 查询工具缺少文本生成器")
}
result, err := t.state.ExecuteTool(ctx, args, generator)
if runtime.Emit != nil {
status := "success"
messageText := "数据库查询完成"
if err != nil {
status = "error"
messageText = "数据库查询失败"
}
runtime.Emit(agents.Frame{Type: "trace", Tool: ToolName, Stage: "execute", Status: status, Message: messageText})
}
return result, err
}
func (t *LoadedTool) RawState() any {
if t == nil {
return nil
}
return t.state
}
func (t *LoadedTool) Close() error {
if t == nil || t.state == nil {
return nil
}
return t.state.Close()
}
type GenerationResult struct {
Database string `json:"database"`
SQL string `json:"sql"`
agents/time/time.go
+94
View File
@@ -1,12 +1,18 @@
package timeagent
import (
"context"
"encoding/json"
"fmt"
"os"
"path/filepath"
"strings"
"time"
agents "aichat/agenttool"
"github.com/volcengine/volcengine-go-sdk/service/arkruntime/model"
"gopkg.in/yaml.v3"
)
const (
@@ -14,10 +20,98 @@ const (
ActivationPrompt = "提供当前日期、时间和常用时间范围。当用户问题包含今天、今日、明天、昨天、本周、本月、本年、最近、历史上的今天、日程安排等相对时间表达时，应先调用此工具；如果后续还需要联网搜索或查数据库，可继续调用 search 或 sql。"
)
type Config struct {
Enabled bool `yaml:"enabled" json:"enabled"`
ActivationPrompt string `yaml:"activation_prompt" json:"activation_prompt"`
}
type ToolArgs struct {
Reason string `json:"reason"`
}
type LoadedTool struct {
cfg *Config
}
func NewLoadedTool(cfg *Config) *LoadedTool {
if cfg == nil {
defaultCfg := defaultConfig()
cfg = &defaultCfg
}
return &LoadedTool{cfg: cfg}
}
func init() {
agents.Register(agents.Descriptor{Name: ToolName, Load: func(path string, options agents.LoadOptions) (agents.LoadedTool, error) {
cfg, err := LoadConfig(path)
if err != nil {
return nil, err
}
return NewLoadedTool(cfg), nil
}})
}
func defaultConfig() Config {
return Config{Enabled: true, ActivationPrompt: ActivationPrompt}
}
func LoadConfig(path string) (*Config, error) {
if _, err := os.Stat(path); err != nil {
if !os.IsNotExist(err) {
return nil, fmt.Errorf("检查时间工具配置失败: %w", err)
}
cfg := defaultConfig()
if err := os.MkdirAll(filepath.Dir(path), 0755); err != nil {
return nil, fmt.Errorf("创建时间工具配置目录失败: %w", err)
}
data, err := yaml.Marshal(&cfg)
if err != nil {
return nil, fmt.Errorf("生成时间工具配置失败: %w", err)
}
if err := os.WriteFile(path, data, 0644); err != nil {
return nil, fmt.Errorf("写入时间工具配置失败: %w", err)
}
}
data, err := os.ReadFile(path)
if err != nil {
return nil, fmt.Errorf("读取时间工具配置失败: %w", err)
}
var cfg Config
if err := yaml.Unmarshal(data, &cfg); err != nil {
return nil, fmt.Errorf("解析时间工具配置失败: %w", err)
}
if strings.TrimSpace(cfg.ActivationPrompt) == "" {
cfg.ActivationPrompt = ActivationPrompt
}
return &cfg, nil
}
func (t *LoadedTool) Name() string { return ToolName }
func (t *LoadedTool) Enabled() bool { return t != nil && t.cfg != nil && t.cfg.Enabled }
func (t *LoadedTool) ToolDefinition(description string) *model.Tool {
if strings.TrimSpace(description) == "" && t != nil && t.cfg != nil {
description = t.cfg.ActivationPrompt
}
return ToolDefinition(description)
}
func (t *LoadedTool) Execute(ctx context.Context, args string, runtime agents.Runtime) (string, error) {
now := runtime.Now
if now.IsZero() {
now = time.Now()
}
result, err := ExecuteTool(args, now)
if err == nil && runtime.Emit != nil {
runtime.Emit(agents.Frame{Type: "trace", Tool: ToolName, Stage: "resolve", Status: "success", Message: "已获取当前时间上下文"})
}
return result, err
}
func (t *LoadedTool) RawState() any { return nil }
type Range struct {
Start time.Time
End time.Time
agenttool/registry.go
+100
View File
@@ -0,0 +1,100 @@
package agenttool
import (
"context"
"fmt"
"sort"
"strings"
"sync"
"time"
"github.com/volcengine/volcengine-go-sdk/service/arkruntime/model"
)
const LegacySearchProfilesKey = "legacy_search_profiles"
type LoadOptions struct {
Values map[string]any
}
func (o LoadOptions) Value(key string) any {
if o.Values == nil {
return nil
}
return o.Values[key]
}
type Frame struct {
Type string
Tool string
Stage string
Status string
Message string
Data map[string]any
Error string
Text string
}
type EmitFunc func(frame any)
type Runtime struct {
Profile any
CompleteText func(context.Context, string, int) (string, error)
Emit EmitFunc
Now time.Time
}
type LoadedTool interface {
Name() string
Enabled() bool
ToolDefinition(description string) *model.Tool
Execute(context.Context, string, Runtime) (string, error)
RawState() any
}
type Descriptor struct {
Name string
Load func(path string, options LoadOptions) (LoadedTool, error)
}
var (
registryMu sync.RWMutex
registry = map[string]Descriptor{}
)
func Register(descriptor Descriptor) {
name := strings.ToLower(strings.TrimSpace(descriptor.Name))
if name == "" {
panic("agenttool: tool name is empty")
}
if descriptor.Load == nil {
panic(fmt.Sprintf("agenttool: %s load function is nil", name))
}
descriptor.Name = name
registryMu.Lock()
defer registryMu.Unlock()
if _, ok := registry[name]; ok {
panic(fmt.Sprintf("agenttool: tool %s already registered", name))
}
registry[name] = descriptor
}
func Lookup(name string) (Descriptor, bool) {
registryMu.RLock()
defer registryMu.RUnlock()
descriptor, ok := registry[strings.ToLower(strings.TrimSpace(name))]
return descriptor, ok
}
func Names() []string {
registryMu.RLock()
defer registryMu.RUnlock()
names := make([]string, 0, len(registry))
for name := range registry {
names = append(names, name)
}
sort.Strings(names)
return names
}
config/config.go
+6 -53
View File
@@ -5,8 +5,6 @@ import (
"os"
"strings"
searchagent "aichat/agents/search"
"gopkg.in/yaml.v3"
)
@@ -17,10 +15,7 @@ const (
defaultToolRouterTimeout = 30
defaultToolRouterMaxTokens = 512
defaultToolRouterSystemText = `你可以按需直接调用可用工具来回答用户问题。
用户询问简单数学计算、四则运算、加减乘除、括号表达式或明确要求算出数值结果时，调用 calculator。
如果用户问题包含今天、今日、明天、昨天、本周、本月、本年、最近等相对时间，且后续需要搜索或查询数据库，应先调用 time 获取绝对日期范围。
需要实时网页资料、新闻、当前版本、近期事件、网页核验或用户明确要求联网时，调用 search。
需要查询本地业务数据、日程、会议、待办、记录、统计或时间范围内数据时，调用 sql。
每个工具的 description 描述了它的适用场景和调用条件。
工具结果优先于模型内置知识；工具失败时必须如实说明，不要编造结果。
只调用确实必要的工具。`
)
@@ -44,7 +39,7 @@ type ToolRouterConfig struct {
Timeout int `yaml:"timeout" json:"timeout"`
MaxTokens int `yaml:"max_tokens" json:"max_tokens"`
SystemPrompt string `yaml:"system_prompt" json:"system_prompt"`
Tools []ToolRouteConfig `yaml:"tools" json:"tools"`
Tools []ToolRouteConfig `yaml:"tools,omitempty" json:"tools,omitempty"`
}
type ToolRouteConfig struct {
@@ -105,12 +100,6 @@ func DefaultToolRouterConfig() ToolRouterConfig {
Timeout: defaultToolRouterTimeout,
MaxTokens: defaultToolRouterMaxTokens,
SystemPrompt: defaultToolRouterSystemText,
Tools: []ToolRouteConfig{
{Name: "calculator", Enabled: true, Description: ""},
{Name: "time", Enabled: true, Description: ""},
{Name: "search", Enabled: true, Description: ""},
{Name: "sql", Enabled: true, Description: ""},
},
}
}
@@ -123,7 +112,7 @@ func Default() Config {
return cfg
}
func Load(path string) (*Config, []searchagent.ProfileConfig, error) {
func Load(path string) (*Config, []map[string]any, error) {
if err := ensureFile(path); err != nil {
return nil, nil, err
}
@@ -285,7 +274,6 @@ func NormalizeToolRouterConfig(cfg *Config) (bool, error) {
func normalizeToolRouterConfig(cfg *Config) (bool, error) {
changed := false
defaults := DefaultToolRouterConfig()
cfg.ToolRouter.OpenAIName = strings.TrimSpace(cfg.ToolRouter.OpenAIName)
if cfg.ToolRouter.Timeout <= 0 {
cfg.ToolRouter.Timeout = defaultToolRouterTimeout
@@ -303,10 +291,6 @@ func normalizeToolRouterConfig(cfg *Config) (bool, error) {
cfg.ToolRouter.SystemPrompt = systemPrompt
changed = true
}
if len(cfg.ToolRouter.Tools) == 0 {
cfg.ToolRouter.Tools = defaults.Tools
changed = true
}
seen := map[string]bool{}
for i := range cfg.ToolRouter.Tools {
tool := &cfg.ToolRouter.Tools[i]
@@ -324,48 +308,17 @@ func normalizeToolRouterConfig(cfg *Config) (bool, error) {
}
seen[name] = true
}
byName := map[string]ToolRouteConfig{}
for _, tool := range cfg.ToolRouter.Tools {
byName[tool.Name] = tool
}
merged := make([]ToolRouteConfig, 0, len(cfg.ToolRouter.Tools)+len(defaults.Tools))
used := map[string]bool{}
for _, tool := range defaults.Tools {
if existing, ok := byName[tool.Name]; ok {
merged = append(merged, existing)
} else {
merged = append(merged, tool)
changed = true
}
used[tool.Name] = true
}
for _, tool := range cfg.ToolRouter.Tools {
if !used[tool.Name] {
merged = append(merged, tool)
}
}
if len(merged) != len(cfg.ToolRouter.Tools) {
changed = true
} else {
for i := range merged {
if merged[i].Name != cfg.ToolRouter.Tools[i].Name {
changed = true
break
}
}
}
cfg.ToolRouter.Tools = merged
return changed, nil
}
func readLegacySearchProfiles(data []byte) []searchagent.ProfileConfig {
func readLegacySearchProfiles(data []byte) []map[string]any {
var legacy struct {
Search searchagent.ProfileConfigs `yaml:"search"`
Search []map[string]any `yaml:"search"`
}
if err := yaml.Unmarshal(data, &legacy); err != nil {
return nil
}
return []searchagent.ProfileConfig(legacy.Search)
return legacy.Search
}
func Write(path string, cfg Config) error {
main.go
+6 -21
View File
@@ -5,12 +5,12 @@ import (
"os"
"strings"
searchagent "aichat/agents/search"
sqlquery "aichat/agents/sql"
agents "aichat/agenttool"
"aichat/config"
"aichat/conversation"
"aichat/llm"
"aichat/server"
"aichat/toolmanager"
"aichat/toolrouter"
)
@@ -27,27 +27,12 @@ func main() {
fmt.Fprintln(os.Stderr, "OpenAI 配置初始化失败:", err)
os.Exit(1)
}
searchConfig, err := searchagent.LoadConfig("agents/search/config.yaml", legacySearchProfiles)
toolManager, err := toolmanager.Load("agents", agents.LoadOptions{Values: map[string]any{agents.LegacySearchProfilesKey: legacySearchProfiles}})
if err != nil {
fmt.Fprintln(os.Stderr, "联网搜索配置加载失败:", err)
fmt.Fprintln(os.Stderr, "工具加载失败:", err)
os.Exit(1)
}
searchState, err := searchagent.NewState(searchConfig)
if err != nil {
fmt.Fprintln(os.Stderr, "联网搜索初始化失败:", err)
os.Exit(1)
}
sqlConfig, err := sqlquery.LoadConfig("agents/sql/config.yaml")
if err != nil {
fmt.Fprintln(os.Stderr, "SQL 查询插件配置加载失败:", err)
os.Exit(1)
}
sqlState, err := sqlquery.NewState(sqlConfig)
if err != nil {
fmt.Fprintln(os.Stderr, "SQL 查询插件初始化失败:", err)
os.Exit(1)
}
defer sqlState.Close()
defer toolManager.Close()
toolRouterState, err := toolrouter.NewState(&cfg.ToolRouter, aiState)
if err != nil {
fmt.Fprintln(os.Stderr, "工具路由配置初始化失败:", err)
@@ -55,7 +40,7 @@ func main() {
}
store := conversation.NewStore("conversations")
app := server.New(cfg, aiState, searchState, sqlState, toolRouterState, store)
app := server.New(cfg, aiState, toolManager, toolRouterState, store)
cfg.Server.Mode = strings.ToLower(cfg.Server.Mode)
if err := app.Run(); err != nil {
fmt.Fprintln(os.Stderr, "服务异常退出:", err)
main_test.go
+23 -13
View File
@@ -7,10 +7,13 @@ import (
"testing"
"time"
"aichat/agents/time"
agents "aichat/agenttool"
"aichat/config"
"aichat/llm"
"aichat/message"
"aichat/stream"
"aichat/toolmanager"
"aichat/toolrouter"
"github.com/volcengine/volcengine-go-sdk/service/arkruntime/model"
@@ -18,6 +21,13 @@ import (
const testOpenAIBaseURL = "https://ark.cn-beijing.volces.com/api/v3"
func newTestToolManager(tools ...agents.LoadedTool) *toolmanager.Manager {
if len(tools) == 0 {
tools = []agents.LoadedTool{timeagent.NewLoadedTool(nil)}
}
return toolmanager.NewForTest(tools...)
}
func newTestAI(t *testing.T, configs []config.OpenAIConfig) *llm.State {
t.Helper()
ai, err := llm.NewState(configs)
@@ -46,7 +56,7 @@ func TestNormalizeToolRouterConfigDefaults(t *testing.T) {
if strings.TrimSpace(cfg.ToolRouter.SystemPrompt) == "" {
t.Fatal("system prompt should be defaulted")
}
if len(cfg.ToolRouter.Tools) != 4 || cfg.ToolRouter.Tools[0].Name != "calculator" || cfg.ToolRouter.Tools[1].Name != "time" || cfg.ToolRouter.Tools[2].Name != "search" || cfg.ToolRouter.Tools[3].Name != "sql" || !cfg.ToolRouter.Tools[0].Enabled || !cfg.ToolRouter.Tools[1].Enabled || !cfg.ToolRouter.Tools[2].Enabled || !cfg.ToolRouter.Tools[3].Enabled {
if len(cfg.ToolRouter.Tools) != 0 {
t.Fatalf("unexpected tools: %#v", cfg.ToolRouter.Tools)
}
}
@@ -66,14 +76,14 @@ func TestNormalizeOpenAIConfigDefaultsContextWindow(t *testing.T) {
}
}
func TestNormalizeToolRouterConfigAddsCalculatorAndTimeBeforeSQL(t *testing.T) {
func TestNormalizeToolRouterConfigKeepsConfiguredTools(t *testing.T) {
cfg := &config.Config{ToolRouter: config.ToolRouterConfig{
Enabled: true,
Timeout: 1,
MaxTokens: 1,
SystemPrompt: "tools",
Tools: []config.ToolRouteConfig{
{Name: "search", Enabled: true},
{Name: " Search ", Enabled: true},
{Name: "sql", Enabled: true},
},
}}
@@ -82,9 +92,9 @@ func TestNormalizeToolRouterConfigAddsCalculatorAndTimeBeforeSQL(t *testing.T) {
t.Fatal(err)
}
if !changed {
t.Fatal("expected calculator and time tools to be added")
t.Fatal("expected configured tools to be normalized")
}
if len(cfg.ToolRouter.Tools) < 4 || cfg.ToolRouter.Tools[0].Name != "calculator" || cfg.ToolRouter.Tools[1].Name != "time" || cfg.ToolRouter.Tools[3].Name != "sql" {
if len(cfg.ToolRouter.Tools) != 2 || cfg.ToolRouter.Tools[0].Name != "search" || cfg.ToolRouter.Tools[1].Name != "sql" {
t.Fatalf("unexpected tool order: %#v", cfg.ToolRouter.Tools)
}
}
@@ -121,15 +131,15 @@ func TestAvailableAgentToolsUsesConfigOrderAndEnabled(t *testing.T) {
t.Fatal(err)
}
tools := toolrouter.AvailableAgentTools(router, ai.ActiveProfile(), nil, nil, nil)
if len(tools) != 2 {
tools := toolrouter.AvailableAgentTools(router, ai.ActiveProfile(), newTestToolManager(), nil)
if len(tools) != 1 {
t.Fatalf("tools length = %d", len(tools))
}
if tools[0].Name() != "calculator" || tools[1].Name() != "time" {
if tools[0].Name() != "time" {
t.Fatalf("unexpected tools: %#v", tools)
}
definition := tools[0].Definition()
if definition.Function == nil || definition.Function.Description != "custom calculator" {
if definition.Function == nil || definition.Function.Description != "custom time" {
t.Fatalf("unexpected definition: %#v", definition)
}
}
@@ -160,7 +170,7 @@ func TestRunAgentToolLoopAppendsToolMessages(t *testing.T) {
t.Fatal(err)
}
messages, err := toolrouter.RunAgentToolLoop(context.Background(), router, ai.ActiveProfile(), []message.ChatMessage{{Role: "user", Content: "今天几号"}}, nil, nil, nil)
messages, err := toolrouter.RunAgentToolLoop(context.Background(), router, ai.ActiveProfile(), []message.ChatMessage{{Role: "user", Content: "今天几号"}}, newTestToolManager(), nil)
if err != nil {
t.Fatal(err)
}
@@ -208,7 +218,7 @@ func TestRunAgentToolLoopMaxIterations(t *testing.T) {
t.Fatal(err)
}
messages, err := toolrouter.RunAgentToolLoop(context.Background(), router, ai.ActiveProfile(), []message.ChatMessage{{Role: "user", Content: "今天"}}, nil, nil, nil)
messages, err := toolrouter.RunAgentToolLoop(context.Background(), router, ai.ActiveProfile(), []message.ChatMessage{{Role: "user", Content: "今天"}}, newTestToolManager(), nil)
if err != nil {
t.Fatal(err)
}
@@ -344,7 +354,7 @@ func TestRunAgentToolLoopImageUsesTextOnlyDecisionMessages(t *testing.T) {
t.Fatal(err)
}
messages, err := toolrouter.RunAgentToolLoop(context.Background(), router, ai.ActiveProfile(), []message.ChatMessage{{Role: "user", Content: "描述这张图", ImageURL: "data:image/png;base64,aGVsbG8="}}, nil, nil, nil)
messages, err := toolrouter.RunAgentToolLoop(context.Background(), router, ai.ActiveProfile(), []message.ChatMessage{{Role: "user", Content: "描述这张图", ImageURL: "data:image/png;base64,aGVsbG8="}}, newTestToolManager(), nil)
if err != nil {
t.Fatal(err)
}
@@ -381,7 +391,7 @@ func TestRunAgentToolLoopUsesConfiguredRouterProfile(t *testing.T) {
t.Fatal(err)
}
_, err = toolrouter.RunAgentToolLoop(context.Background(), router, ai.ActiveProfile(), []message.ChatMessage{{Role: "user", Content: "今天"}}, nil, nil, nil)
_, err = toolrouter.RunAgentToolLoop(context.Background(), router, ai.ActiveProfile(), []message.ChatMessage{{Role: "user", Content: "今天"}}, newTestToolManager(), nil)
if err != nil {
t.Fatal(err)
}
server/handlers.go
+26 -3
View File
@@ -10,6 +10,7 @@ import (
"strings"
"time"
searchagent "aichat/agents/search"
"aichat/contextwindow"
"aichat/conversation"
"aichat/llm"
@@ -61,17 +62,39 @@ func (s *Server) switchOpenAIHandler(c *gin.Context) {
})
}
func (s *Server) searchState() (*searchagent.State, bool) {
if s == nil || s.toolManager == nil {
return nil, false
}
raw, ok := s.toolManager.RawState(searchagent.ToolName)
if !ok {
return nil, false
}
state, ok := raw.(*searchagent.State)
return state, ok && state != nil
}
func (s *Server) listSearchHandler(c *gin.Context) {
c.JSON(http.StatusOK, s.searchState.ListProfiles())
state, ok := s.searchState()
if !ok {
c.JSON(http.StatusOK, searchagent.ListResponse{})
return
}
c.JSON(http.StatusOK, state.ListProfiles())
}
func (s *Server) switchSearchHandler(c *gin.Context) {
state, ok := s.searchState()
if !ok {
c.JSON(http.StatusBadRequest, gin.H{"error": "搜索工具未加载"})
return
}
var req activeProfileRequest
if err := c.ShouldBindJSON(&req); err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": "请求格式错误: " + err.Error()})
return
}
profile, err := s.searchState.SwitchActive(req.Name)
profile, err := state.SwitchActive(req.Name)
if err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
return
@@ -219,7 +242,7 @@ func (s *Server) chatHandler(c *gin.Context) {
contextMessages = chatWindow.ChatMessages
// 用 Function Calling 工具循环替代旧的路由+隐藏上下文机制
messages, err := toolrouter.RunAgentToolLoop(ctx, s.toolRouterState, profile, contextMessages, s.searchState, s.sqlState, emit)
messages, err := toolrouter.RunAgentToolLoop(ctx, s.toolRouterState, profile, contextMessages, s.toolManager, emit)
if err != nil {
fmt.Fprintln(os.Stderr, "Agent 工具循环失败:", err)
messages, err = message.BuildArkMessages(contextMessages)
server/server.go
+4 -7
View File
@@ -6,11 +6,10 @@ import (
"net/http"
"os"
searchagent "aichat/agents/search"
sqlquery "aichat/agents/sql"
"aichat/config"
"aichat/conversation"
"aichat/llm"
"aichat/toolmanager"
"aichat/toolrouter"
"github.com/gin-gonic/gin"
@@ -19,19 +18,17 @@ import (
type Server struct {
cfg *config.Config
aiState *llm.State
searchState *searchagent.State
sqlState *sqlquery.State
toolManager *toolmanager.Manager
toolRouterState *toolrouter.State
store *conversation.Store
router *gin.Engine
}
func New(cfg *config.Config, aiState *llm.State, searchState *searchagent.State, sqlState *sqlquery.State, toolRouterState *toolrouter.State, store *conversation.Store) *Server {
func New(cfg *config.Config, aiState *llm.State, toolManager *toolmanager.Manager, toolRouterState *toolrouter.State, store *conversation.Store) *Server {
s := &Server{
cfg: cfg,
aiState: aiState,
searchState: searchState,
sqlState: sqlState,
toolManager: toolManager,
toolRouterState: toolRouterState,
store: store,
router: gin.Default(),
toolmanager/manager.go
+125
View File
@@ -0,0 +1,125 @@
package toolmanager
import (
"errors"
"fmt"
"os"
"path/filepath"
"sort"
"strings"
_ "aichat/agents/calculator"
_ "aichat/agents/search"
_ "aichat/agents/sql"
_ "aichat/agents/time"
agents "aichat/agenttool"
)
type Manager struct {
tools map[string]agents.LoadedTool
order []string
}
func Load(root string, options agents.LoadOptions) (*Manager, error) {
entries, err := os.ReadDir(root)
if err != nil {
return nil, fmt.Errorf("扫描工具目录失败: %w", err)
}
manager := &Manager{tools: map[string]agents.LoadedTool{}}
for _, entry := range entries {
if !entry.IsDir() {
continue
}
name := strings.ToLower(strings.TrimSpace(entry.Name()))
descriptor, ok := agents.Lookup(name)
if !ok {
continue
}
tool, err := descriptor.Load(filepath.Join(root, entry.Name(), "config.yaml"), options)
if err != nil {
manager.Close()
return nil, fmt.Errorf("加载工具 %s 失败: %w", name, err)
}
if tool == nil {
continue
}
toolName := strings.ToLower(strings.TrimSpace(tool.Name()))
if toolName == "" {
toolName = name
}
if _, ok := manager.tools[toolName]; ok {
manager.Close()
return nil, fmt.Errorf("工具名称重复: %s", toolName)
}
manager.tools[toolName] = tool
manager.order = append(manager.order, toolName)
}
return manager, nil
}
func NewForTest(tools ...agents.LoadedTool) *Manager {
manager := &Manager{tools: map[string]agents.LoadedTool{}}
for _, tool := range tools {
if tool == nil {
continue
}
name := strings.ToLower(strings.TrimSpace(tool.Name()))
if name == "" {
continue
}
if _, ok := manager.tools[name]; !ok {
manager.order = append(manager.order, name)
}
manager.tools[name] = tool
}
return manager
}
func (m *Manager) Tools() []agents.LoadedTool {
if m == nil {
return nil
}
tools := make([]agents.LoadedTool, 0, len(m.order))
for _, name := range m.order {
if tool := m.tools[name]; tool != nil {
tools = append(tools, tool)
}
}
return tools
}
func (m *Manager) Get(name string) (agents.LoadedTool, bool) {
if m == nil {
return nil, false
}
tool, ok := m.tools[strings.ToLower(strings.TrimSpace(name))]
return tool, ok
}
func (m *Manager) RawState(name string) (any, bool) {
tool, ok := m.Get(name)
if !ok || tool == nil {
return nil, false
}
return tool.RawState(), true
}
func (m *Manager) Close() error {
if m == nil {
return nil
}
var errs []string
for _, tool := range m.Tools() {
if closer, ok := tool.(interface{ Close() error }); ok {
if err := closer.Close(); err != nil {
errs = append(errs, err.Error())
}
}
}
if len(errs) > 0 {
sort.Strings(errs)
return errors.New(strings.Join(errs, "; "))
}
return nil
}
toolrouter/loop.go
+3 -4
View File
@@ -6,11 +6,10 @@ import (
"strings"
"time"
searchagent "aichat/agents/search"
sqlquery "aichat/agents/sql"
"aichat/llm"
"aichat/message"
"aichat/stream"
"aichat/toolmanager"
"aichat/utils"
"github.com/volcengine/volcengine-go-sdk/service/arkruntime/model"
@@ -18,7 +17,7 @@ import (
const maxAgentToolIterations = 6
func RunAgentToolLoop(ctx context.Context, state *State, profile *llm.Profile, chatMessages []message.ChatMessage, searchState *searchagent.State, sqlState *sqlquery.State, emit stream.EmitFunc) ([]*model.ChatCompletionMessage, error) {
func RunAgentToolLoop(ctx context.Context, state *State, profile *llm.Profile, chatMessages []message.ChatMessage, manager *toolmanager.Manager, emit stream.EmitFunc) ([]*model.ChatCompletionMessage, error) {
finalMessages, err := message.BuildArkMessages(chatMessages)
if err != nil {
return nil, err
@@ -27,7 +26,7 @@ func RunAgentToolLoop(ctx context.Context, state *State, profile *llm.Profile, c
if state != nil {
routerProfile = state.RouterProfile(profile)
}
tools := availableAgentTools(state, routerProfile, searchState, sqlState, emit)
tools := availableAgentTools(state, routerProfile, manager, emit)
if len(tools) == 0 {
return finalMessages, nil
}
toolrouter/tools.go
+81 -87
View File
@@ -6,14 +6,13 @@ import (
"strings"
"time"
"aichat/agents/calculator"
searchagent "aichat/agents/search"
sqlquery "aichat/agents/sql"
timeagent "aichat/agents/time"
agents "aichat/agenttool"
"aichat/completion"
"aichat/config"
"aichat/llm"
"aichat/message"
"aichat/stream"
"aichat/toolmanager"
"aichat/utils"
"github.com/volcengine/volcengine-go-sdk/service/arkruntime/model"
@@ -33,101 +32,96 @@ func (t AgentTool) Name() string { return t.name }
func (t AgentTool) Definition() *model.Tool { return t.definition }
func AvailableAgentTools(state *State, profile *llm.Profile, searchState *searchagent.State, sqlState *sqlquery.State, emit stream.EmitFunc) []AgentTool {
return availableAgentTools(state, profile, searchState, sqlState, emit)
func AvailableAgentTools(state *State, profile *llm.Profile, manager *toolmanager.Manager, emit stream.EmitFunc) []AgentTool {
return availableAgentTools(state, profile, manager, emit)
}
func availableAgentTools(state *State, profile *llm.Profile, searchState *searchagent.State, sqlState *sqlquery.State, emit stream.EmitFunc) []AgentTool {
if state == nil || state.cfg == nil || !state.cfg.Enabled {
func availableAgentTools(state *State, profile *llm.Profile, manager *toolmanager.Manager, emit stream.EmitFunc) []AgentTool {
if state == nil || state.cfg == nil || !state.cfg.Enabled || manager == nil {
return nil
}
tools := make([]AgentTool, 0, len(state.cfg.Tools))
overrides := map[string]config.ToolRouteConfig{}
for _, item := range state.cfg.Tools {
if !item.Enabled {
name := strings.ToLower(strings.TrimSpace(item.Name))
if name != "" {
overrides[name] = item
}
}
if len(state.cfg.Tools) > 0 {
tools := make([]AgentTool, 0, len(state.cfg.Tools))
for _, item := range state.cfg.Tools {
name := strings.ToLower(strings.TrimSpace(item.Name))
if name == "" || !item.Enabled {
continue
}
tool, ok := manager.Get(name)
if !ok {
continue
}
if agentTool, ok := buildAgentTool(tool, strings.TrimSpace(item.Description), profile, emit); ok {
tools = append(tools, agentTool)
}
}
return tools
}
loaded := manager.Tools()
tools := make([]AgentTool, 0, len(loaded))
for _, tool := range loaded {
if tool == nil {
continue
}
description := strings.TrimSpace(item.Description)
switch item.Name {
case calculator.ToolName:
tools = append(tools, AgentTool{
name: calculator.ToolName,
definition: calculator.ToolDefinition(description),
execute: func(ctx context.Context, args string) (string, error) {
result, err := calculator.ExecuteTool(args)
if err == nil && emit != nil {
emit(stream.Frame{Type: "trace", Tool: calculator.ToolName, Stage: "calculate", Status: "success", Message: "四则运算完成"})
}
return result, err
},
})
case timeagent.ToolName:
tools = append(tools, AgentTool{
name: timeagent.ToolName,
definition: timeagent.ToolDefinition(description),
execute: func(ctx context.Context, args string) (string, error) {
result, err := timeagent.ExecuteTool(args, time.Now())
if err == nil && emit != nil {
emit(stream.Frame{Type: "trace", Tool: timeagent.ToolName, Stage: "resolve", Status: "success", Message: "已获取当前时间上下文"})
}
return result, err
},
})
case searchagent.ToolName:
if searchState == nil || !searchState.Enabled() {
continue
}
tools = append(tools, AgentTool{
name: searchagent.ToolName,
definition: searchState.ToolDefinition(description),
execute: func(ctx context.Context, args string) (string, error) {
if emit != nil {
emit(stream.Frame{Type: "trace", Tool: searchagent.ToolName, Stage: "request", Status: "running", Message: "正在联网搜索"})
}
result, err := searchState.ExecuteTool(ctx, args)
if emit != nil {
status := "success"
messageText := "联网搜索完成"
if err != nil {
status = "error"
messageText = "联网搜索失败"
}
emit(stream.Frame{Type: "trace", Tool: searchagent.ToolName, Stage: "results", Status: status, Message: messageText})
}
return result, err
},
})
case sqlquery.ToolName:
if sqlState == nil || !sqlState.Enabled() {
continue
}
tools = append(tools, AgentTool{
name: sqlquery.ToolName,
definition: sqlState.ToolDefinition(description),
execute: func(ctx context.Context, args string) (string, error) {
if emit != nil {
emit(stream.Frame{Type: "trace", Tool: sqlquery.ToolName, Stage: "execute", Status: "running", Message: "正在查询数据库"})
}
generator := func(ctx context.Context, prompt string, maxTokens int) (string, error) {
return completion.CompleteText(ctx, profile, []message.ChatMessage{{Role: "system", Content: prompt}}, maxTokens)
}
result, err := sqlState.ExecuteTool(ctx, args, generator)
if emit != nil {
status := "success"
messageText := "数据库查询完成"
if err != nil {
status = "error"
messageText = "数据库查询失败"
}
emit(stream.Frame{Type: "trace", Tool: sqlquery.ToolName, Stage: "execute", Status: status, Message: messageText})
}
return result, err
},
})
item := overrides[tool.Name()]
if agentTool, ok := buildAgentTool(tool, strings.TrimSpace(item.Description), profile, emit); ok {
tools = append(tools, agentTool)
}
}
return tools
}
func buildAgentTool(tool agents.LoadedTool, description string, profile *llm.Profile, emit stream.EmitFunc) (AgentTool, bool) {
if tool == nil || !tool.Enabled() {
return AgentTool{}, false
}
definition := tool.ToolDefinition(description)
if definition == nil || definition.Function == nil {
return AgentTool{}, false
}
name := tool.Name()
return AgentTool{
name: name,
definition: definition,
execute: func(ctx context.Context, args string) (string, error) {
runtime := agents.Runtime{
Profile: profile,
Now: time.Now(),
Emit: wrapAgentEmit(emit),
}
if profile != nil {
runtime.CompleteText = func(ctx context.Context, prompt string, maxTokens int) (string, error) {
return completion.CompleteText(ctx, profile, []message.ChatMessage{{Role: "system", Content: prompt}}, maxTokens)
}
}
return tool.Execute(ctx, args, runtime)
},
}, true
}
func wrapAgentEmit(emit stream.EmitFunc) agents.EmitFunc {
if emit == nil {
return nil
}
return func(frame any) {
switch value := frame.(type) {
case agents.Frame:
emit(stream.Frame{Type: value.Type, Tool: value.Tool, Stage: value.Stage, Status: value.Status, Message: value.Message, Data: value.Data, Error: value.Error, Text: value.Text})
case stream.Frame:
emit(value)
}
}
}
func ExecuteAgentToolCall(ctx context.Context, call *model.ToolCall, tools map[string]AgentTool, emit stream.EmitFunc) string {
if call == nil || call.Type != model.ToolTypeFunction {
result := "工具调用无效：仅支持 function 类型工具。"