agents_and_robots/tools/registry.go

package tools

import (
	"context"
	"encoding/json"
	"fmt"
	"log/slog"
	"sort"
	"time"

	coretypes "github.com/enmanuel/agents/pkg/llm"
	"github.com/enmanuel/agents/shell/logger"
)

// AuditFunc is called after each tool execution for audit purposes.
// The registry does not depend on the audit package directly.
type AuditFunc func(toolName string, durationMS int64, err error)

// Registry holds available tools keyed by name.
type Registry struct {
	tools       map[string]Tool
	logger      *slog.Logger
	rateLimiter *RateLimiter // nil when rate limiting is disabled
	auditFn     AuditFunc   // nil when audit is disabled
}

// NewRegistry creates an empty registry.
func NewRegistry(log *slog.Logger) *Registry {
	return &Registry{
		tools:  make(map[string]Tool),
		logger: log.With(logger.FieldComponent, "tools"),
	}
}

// Register adds a tool to the registry.
func (r *Registry) Register(t Tool) {
	r.tools[t.Def.Name] = t
	r.logger.Debug("tool_registered", "name", t.Def.Name)
}

// Get looks up a tool by name.
func (r *Registry) Get(name string) (Tool, bool) {
	t, ok := r.tools[name]
	return t, ok
}

// Names returns all registered tool names in sorted order.
func (r *Registry) Names() []string {
	names := make([]string, 0, len(r.tools))
	for k := range r.tools {
		names = append(names, k)
	}
	sort.Strings(names)
	return names
}

// Len returns the number of registered tools.
func (r *Registry) Len() int {
	return len(r.tools)
}

// SetRateLimiter attaches a rate limiter to the registry.
// When set, ExecuteForRoom checks the limit before running the tool.
func (r *Registry) SetRateLimiter(rl *RateLimiter) {
	r.rateLimiter = rl
}

// SetAuditFunc attaches an audit callback to the registry.
// When set, it is called after each tool execution.
func (r *Registry) SetAuditFunc(fn AuditFunc) {
	r.auditFn = fn
}

// ExecuteForRoom is like Execute but checks the per-room rate limit first.
// If the rate limit is exceeded, it returns an error result without executing.
func (r *Registry) ExecuteForRoom(ctx context.Context, name, argsJSON, roomID string) Result {
	if r.rateLimiter != nil && roomID != "" {
		if !r.rateLimiter.Allow(roomID) {
			r.logger.Warn("tool_rate_limited", "tool", name, "room", roomID)
			return Result{Err: fmt.Errorf("rate limit exceeded for room %s: too many tool calls per minute", roomID)}
		}
	}
	return r.Execute(ctx, name, argsJSON)
}

// Execute looks up a tool by name and runs it. Returns an error result if not found.
func (r *Registry) Execute(ctx context.Context, name string, argsJSON string) Result {
	t, ok := r.tools[name]
	if !ok {
		r.logger.Warn("tool_not_found", "tool", name)
		return Result{Err: fmt.Errorf("tool %q not found", name)}
	}

	var args map[string]any
	if argsJSON != "" {
		if err := json.Unmarshal([]byte(argsJSON), &args); err != nil {
			r.logger.Warn("tool_args_invalid", "tool", name, "err", err)
			return Result{Err: fmt.Errorf("parse args for %q: %w", name, err)}
		}
	}

	r.logger.Info("tool_exec_start", "tool", name)
	start := time.Now()
	result := t.Exec(ctx, args)
	ms := time.Since(start).Milliseconds()

	if result.Err != nil {
		r.logger.Warn("tool_exec_error", "tool", name, "err", result.Err, logger.FieldDurationMS, ms)
	} else {
		r.logger.Info("tool_exec_end", "tool", name, logger.FieldDurationMS, ms)
	}

	// Audit callback
	if r.auditFn != nil {
		r.auditFn(name, ms, result.Err)
	}

	return result
}

// ToLLMSpecs converts all registered tools to the LLM-compatible ToolSpec format.
// This is a pure transformation — no side effects.
func (r *Registry) ToLLMSpecs() []coretypes.ToolSpec {
	specs := make([]coretypes.ToolSpec, 0, len(r.tools))
	for _, name := range r.Names() {
		t := r.tools[name]
		specs = append(specs, defToLLMSpec(t.Def))
	}
	return specs
}

// defToLLMSpec converts a pure Def to an LLM ToolSpec with JSON Schema.
func defToLLMSpec(d Def) coretypes.ToolSpec {
	properties := make(map[string]any, len(d.Parameters))
	required := make([]string, 0)

	for _, p := range d.Parameters {
		properties[p.Name] = map[string]any{
			"type":        p.Type,
			"description": p.Description,
		}
		if p.Required {
			required = append(required, p.Name)
		}
	}

	schema := map[string]any{
		"type":       "object",
		"properties": properties,
	}
	if len(required) > 0 {
		schema["required"] = required
	}

	return coretypes.ToolSpec{
		Name:        d.Name,
		Description: d.Description,
		InputSchema: schema,
	}
}