raven/internal/tool/tool.go

// Package tool provides a registry for external tools that can be invoked by
// LLMs.
//
// The package bridges YAML configuration to exec.CommandContext, allowing
// tools to be defined declaratively without writing Go code. Each tool
// specifies a command, argument templates using Go's text/template syntax,
// JSON Schema parameters for LLM input, and an optional execution timeout.
//
// The registry validates all tool definitions at construction time,
// failing fast on configuration errors. At execution time, it expands
// argument templates with LLM-provided JSON, runs the subprocess, and
// returns stdout.
package tool

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"os/exec"
	"text/template"
	"time"

	"github.com/sashabaranov/go-openai"
)

// fn provides template functions available to argument templates.
var fn = template.FuncMap{
	"json": func(v any) string {
		b, _ := json.Marshal(v)
		return string(b)
	},
}

// Tool represents an external tool that can be invoked by LLMs.
// Tools are executed as subprocesses with templated arguments.
type Tool struct {
	// Name uniquely identifies the tool within a registry.
	Name string `yaml:"name"`
	// Description explains the tool's purpose for the LLM.
	Description string `yaml:"description"`
	// Command is the executable path or name.
	Command string `yaml:"command"`
	// Arguments are Go templates expanded with LLM-provided parameters.
	// Empty results after expansion are filtered out.
	Arguments []string `yaml:"arguments"`
	// Parameters is a JSON Schema describing expected input from the LLM.
	Parameters map[string]any `yaml:"parameters"`
	// Timeout limits execution time (e.g., "30s", "5m").
	// Empty means no timeout.
	Timeout string `yaml:"timeout"`
	// timeout is the parsed duration, set during registry construction.
	timeout time.Duration `yaml:"-"`
}

// OpenAI converts the tool to an OpenAI function definition for API calls.
func (t Tool) OpenAI() openai.Tool {
	return openai.Tool{
		Type: openai.ToolTypeFunction,
		Function: &openai.FunctionDefinition{
			Name:        t.Name,
			Description: t.Description,
			Parameters:  t.Parameters,
		},
	}
}

// ParseArguments expands argument templates with the provided JSON data.
// The args parameter should be a JSON object string; empty string or "{}"
// results in an empty data map. Templates producing empty strings are
// filtered from the result, allowing conditional arguments.
func (t Tool) ParseArguments(args string) ([]string, error) {
	var data = map[string]any{}
	if args != "" && args != "{}" {
		if err := json.Unmarshal([]byte(args), &data); err != nil {
			return nil, fmt.Errorf(
				"invalid arguments JSON: %w", err,
			)
		}
	}

	var result []string
	for _, v := range t.Arguments {
		tmpl, err := template.New("").Funcs(fn).Parse(v)
		if err != nil {
			return nil, fmt.Errorf(
				"invalid template %q: %w", v, err,
			)
		}

		var buf bytes.Buffer
		if err := tmpl.Execute(&buf, data); err != nil {
			return nil, fmt.Errorf(
				"execute template %q: %w", v, err,
			)
		}

		// Filter out empty strings (unused conditional arguments).
		if s := buf.String(); s != "" {
			result = append(result, s)
		}
	}

	return result, nil
}

// Validate checks that the tool definition is complete and valid.
// It verifies required fields are present, the timeout (if specified)
// is parseable, and all argument templates are syntactically valid.
func (t Tool) Validate() error {
	if t.Name == "" {
		return fmt.Errorf("missing name")
	}
	if t.Description == "" {
		return fmt.Errorf("missing description")
	}
	if t.Command == "" {
		return fmt.Errorf("missing command")
	}
	if t.Timeout != "" {
		if _, err := time.ParseDuration(t.Timeout); err != nil {
			return fmt.Errorf("invalid timeout: %v", err)
		}
	}
	for _, arg := range t.Arguments {
		if _, err := template.New("").Funcs(fn).Parse(arg); err != nil {
			return fmt.Errorf("invalid argument template")
		}
	}

	return nil
}

// Registry holds tools indexed by name and handles their execution.
// It validates all tools at construction time to fail fast on
// configuration errors.
type Registry struct {
	tools map[string]*Tool
}

// NewRegistry creates a registry from the provided tool definitions.
// Returns an error if any tool fails validation or if duplicate names exist.
func NewRegistry(tools []Tool) (*Registry, error) {
	var r = &Registry{
		tools: make(map[string]*Tool),
	}
	for _, t := range tools {
		if err := t.Validate(); err != nil {
			return nil, fmt.Errorf("invalid tool: %v", err)
		}
		if t.Timeout != "" {
			d, err := time.ParseDuration(t.Timeout)
			if err != nil {
				return nil, fmt.Errorf(
					"parse timeout: %v", err,
				)
			}
			t.timeout = d
		}

		if _, exists := r.tools[t.Name]; exists {
			return nil, fmt.Errorf(
				"duplicate tool name: %s", t.Name,
			)
		}
		r.tools[t.Name] = &t
	}

	return r, nil
}

// Get returns a tool by name and a boolean indicating if it was found.
func (r *Registry) Get(name string) (*Tool, bool) {
	tool, ok := r.tools[name]
	return tool, ok
}

// List returns all registered tool names in arbitrary order.
func (r *Registry) List() []string {
	names := make([]string, 0, len(r.tools))
	for name := range r.tools {
		names = append(names, name)
	}
	return names
}

// Execute runs a tool by name with the provided JSON arguments.
// It expands argument templates, executes the command as a subprocess,
// and returns stdout on success. The context can be used for cancellation;
// tool-specific timeouts are applied on top of any context deadline.
func (r *Registry) Execute(
	ctx context.Context, name string, args string,
) (string, error) {
	tool, ok := r.tools[name]
	if !ok {
		return "", fmt.Errorf("unknown tool: %s", name)
	}

	// Evaluate argument templates.
	cmdArgs, err := tool.ParseArguments(args)
	if err != nil {
		return "", fmt.Errorf("parse arguments: %w", err)
	}

	// If defined, use the timeout.
	if tool.timeout > 0 {
		var cancel context.CancelFunc
		ctx, cancel = context.WithTimeout(ctx, tool.timeout)
		defer cancel()
	}

	// Setup and run the command.
	var (
		stdout, stderr bytes.Buffer
		cmd            = exec.CommandContext(
			ctx, tool.Command, cmdArgs...,
		)
	)
	cmd.Stdout = &stdout
	cmd.Stderr = &stderr
	if err := cmd.Run(); err != nil {
		if ctx.Err() == context.DeadlineExceeded && tool.timeout > 0 {
			return "", fmt.Errorf(
				"tool %s timed out after %v",
				name, tool.timeout,
			)
		}
		return "", fmt.Errorf(
			"tool %s: %w\nstderr: %s",
			name, err, stderr.String(),
		)
	}

	return stdout.String(), nil
}