odidere/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"

	"code.chimeric.al/chimerical/odidere/internal/config"
)

// 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
	// Description explains the tool's purpose for the LLM.
	Description string
	// Command is the executable path or name.
	Command string
	// Arguments are Go templates expanded with LLM-provided parameters.
	Arguments []string
	// Parameters is a JSON Schema describing expected input from the LLM.
	Parameters map[string]any
	// timeout is the parsed execution duration limit.
	timeout time.Duration
}

// NewTool creates a Tool from a ToolConfig.
func NewTool(cfg config.ToolConfig) (*Tool, error) {
	var timeout time.Duration
	if cfg.Timeout != "" {
		var err error
		timeout, err = time.ParseDuration(cfg.Timeout)
		if err != nil {
			return nil, fmt.Errorf(
				"invalid timeout %q: %w", cfg.Timeout, err,
			)
		}
	}

	return &Tool{
		Name:        cfg.Name,
		Description: cfg.Description,
		Command:     cfg.Command,
		Arguments:   cfg.Arguments,
		Parameters:  cfg.Parameters,
		timeout:     timeout,
	}, nil
}

// 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
}

// Registry holds tools indexed by name and handles their execution.
// It validates all tool definitions 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 []config.ToolConfig) (*Registry, error) {
	var r = &Registry{
		tools: make(map[string]*Tool),
	}
	for _, tc := range tools {
		t, err := NewTool(tc)
		if err != nil {
			return nil, fmt.Errorf("invalid tool %q: %w", tc.Name, err)
		}

		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
}