odidere/internal/llm/tool.go

package llm

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

	"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,
// and described to the LLM via their JSON Schema parameters.
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
}

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

// ToAPI converts the tool definition into the API format suitable for
// inclusion in a ChatRequest.
func (t *Tool) ToAPI() APITool {
	return APITool{
		Type: "function",
		Function: &FunctionDef{
			Name:        t.Name,
			Description: t.Description,
			Parameters:  t.Parameters,
		},
	}
}

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

// ToAPI converts all registered tools into API format suitable for
// inclusion in a ChatRequest.
func (r *Registry) ToAPI() []APITool {
	if r == nil {
		return nil
	}
	tools := make([]APITool, 0, len(r.tools))
	for _, t := range r.tools {
		tools = append(tools, t.ToAPI())
	}
	return tools
}

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