102 lines
2.7 KiB
Go
102 lines
2.7 KiB
Go
package openapi
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
)
|
|
|
|
// ToolHandler is a function that implements a tool's logic.
|
|
type ToolHandler func(ctx context.Context, s *Session, args json.RawMessage) (any, error)
|
|
|
|
// ToolDefinition pairs the OpenAI-compatible tool schema with a Go handler.
|
|
type ToolDefinition struct {
|
|
Tool Tool
|
|
Handler ToolHandler
|
|
}
|
|
|
|
// ToolRegistry defines the interface for tool discovery and execution.
|
|
type ToolRegistry interface {
|
|
Tools(ctx context.Context) []Tool
|
|
Execute(ctx context.Context, s *Session, name string, args json.RawMessage) (any, error)
|
|
}
|
|
|
|
// MapRegistry is a simple map-based implementation of ToolRegistry.
|
|
type MapRegistry struct {
|
|
tools map[string]ToolDefinition
|
|
}
|
|
|
|
// NewMapRegistry creates a new MapRegistry.
|
|
func NewMapRegistry() *MapRegistry {
|
|
return &MapRegistry{tools: make(map[string]ToolDefinition)}
|
|
}
|
|
|
|
// Register adds a tool to the registry.
|
|
func (r *MapRegistry) Register(name, description string, parameters any, handler ToolHandler) error {
|
|
paramsJSON, err := json.Marshal(parameters)
|
|
if err != nil {
|
|
return fmt.Errorf("marshal parameters: %w", err)
|
|
}
|
|
|
|
r.tools[name] = ToolDefinition{
|
|
Tool: Tool{
|
|
Type: "function",
|
|
Function: FunctionDefinition{
|
|
Name: name,
|
|
Description: description,
|
|
Parameters: paramsJSON,
|
|
},
|
|
},
|
|
Handler: handler,
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// Tools returns the list of tools in the registry.
|
|
func (r *MapRegistry) Tools(ctx context.Context) []Tool {
|
|
tools := make([]Tool, 0, len(r.tools))
|
|
for _, td := range r.tools {
|
|
tools = append(tools, td.Tool)
|
|
}
|
|
return tools
|
|
}
|
|
|
|
// Execute looks up and runs a tool by name.
|
|
func (r *MapRegistry) Execute(ctx context.Context, s *Session, name string, args json.RawMessage) (any, error) {
|
|
td, ok := r.tools[name]
|
|
if !ok {
|
|
return nil, fmt.Errorf("tool %s not found", name)
|
|
}
|
|
return td.Handler(ctx, s, args)
|
|
}
|
|
|
|
// CompositeRegistry allows combining multiple tool registries.
|
|
type CompositeRegistry struct {
|
|
registries []ToolRegistry
|
|
}
|
|
|
|
// Add adds a registry to the composite.
|
|
func (r *CompositeRegistry) Add(registry ToolRegistry) {
|
|
r.registries = append(r.registries, registry)
|
|
}
|
|
|
|
// Tools returns the union of all tools from all registries.
|
|
func (r *CompositeRegistry) Tools(ctx context.Context) []Tool {
|
|
var all []Tool
|
|
for _, reg := range r.registries {
|
|
all = append(all, reg.Tools(ctx)...)
|
|
}
|
|
return all
|
|
}
|
|
|
|
// Execute attempts to execute the tool by trying each registry in order.
|
|
func (r *CompositeRegistry) Execute(ctx context.Context, s *Session, name string, args json.RawMessage) (any, error) {
|
|
for _, reg := range r.registries {
|
|
res, err := reg.Execute(ctx, s, name, args)
|
|
if err == nil {
|
|
return res, nil
|
|
}
|
|
}
|
|
return nil, fmt.Errorf("tool %s not found in any registry", name)
|
|
}
|