126 lines
2.7 KiB
Go
126 lines
2.7 KiB
Go
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
|
|
}
|