lmcli/pkg/cli/functions.go

package cli

import (
	"database/sql"
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"strings"

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

type FunctionResult struct {
	Message string `json:"message"`
	Result  any    `json:"result,omitempty"`
}

type FunctionParameter struct {
	Type        string   `json:"type"` // "string", "integer", "boolean"
	Description string   `json:"description"`
	Enum        []string `json:"enum,omitempty"`
}

type FunctionParameters struct {
	Type       string                       `json:"type"` // "object"
	Properties map[string]FunctionParameter `json:"properties"`
	Required   []string                     `json:"required,omitempty"` // required function parameter names
}

type AvailableTool struct {
	openai.Tool
	// The tool's implementation. Returns a string, as tool call results
	// are treated as normal messages with string contents.
	Impl func(arguments map[string]interface{}) (string, error)
}

var AvailableTools = map[string]AvailableTool{
	"read_dir": {
		Tool: openai.Tool{Type: "function", Function: openai.FunctionDefinition{
			Name: "read_dir",
			Description: `Return the contents of the CWD (current working directory).

Results are returned as JSON in the following format:
{
	"message": "success", // if successful, or a different message indicating failure
	"result": [
		{"name": "a_file", "type": "file", "size": 123},
		{"name": "a_directory/", "type": "dir", "size": 11},
		... // more files or directories
	]
}

For files, size represents the size (in bytes) of the file.
For directories, size represents the number of entries in that directory.`,
			Parameters: FunctionParameters{
				Type: "object",
				Properties: map[string]FunctionParameter{
					"relative_dir": {
						Type:        "string",
						Description: "If set, read the contents of a directory relative to the current one.",
					},
				},
			},
		}},
		Impl: func(args map[string]interface{}) (string, error) {
			var relativeDir string
			tmp, ok := args["relative_dir"]
			if ok {
				relativeDir, ok = tmp.(string)
				if !ok {
					return "", fmt.Errorf("Invalid relative_dir in function arguments: %v", tmp)
				}
			}
			return ReadDir(relativeDir), nil
		},
	},
}

func resultToJson(result FunctionResult) string {
	if result.Message == "" {
		// When message not supplied, assume success
		result.Message = "success"
	}

	jsonBytes, err := json.Marshal(result)
	if err != nil {
		fmt.Printf("Could not marshal FunctionResult to JSON: %v\n", err)
	}
	return string(jsonBytes)
}

// ExecuteToolCalls handles the execution of all tool_calls provided, and
// returns their results formatted as []Message(s) with role: 'tool' and.
func ExecuteToolCalls(toolCalls []openai.ToolCall) ([]Message, error) {
	var toolResults []Message
	for _, toolCall := range toolCalls {
		if toolCall.Type != "function" {
			// unsupported tool type
			continue
		}

		tool, ok := AvailableTools[toolCall.Function.Name]
		if !ok {
			return nil, fmt.Errorf("Requested tool '%s' does not exist. Hallucination?", toolCall.Function.Name)
		}

		var functionArgs map[string]interface{}
		err := json.Unmarshal([]byte(toolCall.Function.Arguments), &functionArgs)
		if err != nil {
			return nil, fmt.Errorf("Could not unmarshal tool arguments. Malformed JSON? Error: %v", err)
		}

		// TODO: ability to silence this
		fmt.Fprintf(os.Stderr, "INFO: Executing tool '%s' with args %s\n", toolCall.Function.Name, toolCall.Function.Arguments)

		// Execute the tool
		toolResult, err := tool.Impl(functionArgs)
		if err != nil {
			// This can happen if the model missed or supplied invalid tool args
			return nil, fmt.Errorf("Tool '%s' error: %v\n", toolCall.Function.Name, err)
		}

		toolResults = append(toolResults, Message{
			Role:            "tool",
			OriginalContent: toolResult,
			ToolCallID:      sql.NullString{String: toolCall.ID, Valid: true},
			// name is not required since the introduction of ToolCallID
			// hypothesis: by setting it, we inform the model of what a
			// function's purpose was if future requests omit the function
			// definition
		})
	}
	return toolResults, nil
}

// isPathContained attempts to verify whether `path` is the same as or
// contained within `directory`. It is overly cautious, returning false even if
// `path` IS contained within `directory`, but the two paths use different
// casing, and we happen to be on a case-insensitive filesystem.
// This is ultimately to attempt to stop an LLM from going outside of where I
// tell it to. Additional layers of security should be considered.. run in a
// VM/container.
func isPathContained(directory string, path string) (bool, error) {
	// Clean and resolve symlinks for both paths
	absPath, err := filepath.Abs(path)
	if err != nil {
		return false, err
	}
	realPath, err := filepath.EvalSymlinks(absPath)
	if err != nil {
		return false, err
	}

	absDirectory, err := filepath.Abs(directory)
	if err != nil {
		return false, err
	}
	realDirectory, err := filepath.EvalSymlinks(absDirectory)
	if err != nil {
		return false, err
	}

	// Case insensitive checks
	if !strings.EqualFold(realPath, realDirectory) &&
		!strings.HasPrefix(strings.ToLower(realPath), strings.ToLower(realDirectory)+string(os.PathSeparator)) {
		return false, nil
	}

	return true, nil
}

func isPathWithinCWD(path string) (bool, *FunctionResult) {
	cwd, err := os.Getwd()
	if err != nil {
		return false, &FunctionResult{Message: "Failed to determine current working directory"}
	}
	if ok, err := isPathContained(cwd, path); !ok {
		if err != nil {
			return false, &FunctionResult{Message: fmt.Sprintf("Could not determine whether path '%s' is within the current working directory: %s", path, err.Error())}
		}
		return false, &FunctionResult{Message: fmt.Sprintf("Path '%s' is not within the current working directory", path)}
	}
	return true, nil
}

func ReadDir(path string) string {
	// TODO(?): implement whitelist - list of directories which model is allowed to work in
	if path == "" {
		path = "."
	}
	ok, res := isPathWithinCWD(path)
	if !ok {
		return resultToJson(*res)
	}

	files, err := os.ReadDir(path)
	if err != nil {
		return resultToJson(FunctionResult{
			Message: err.Error(),
		})
	}

	var dirContents []map[string]interface{}
	for _, f := range files {
		info, _ := f.Info()

		name := f.Name()
		entryType := "file"
		size := info.Size()

		if info.IsDir() {
			name += "/"
			entryType = "dir"
			subdirfiles, _ := os.ReadDir(filepath.Join(".", path, info.Name()))
			size = int64(len(subdirfiles))
		}

		dirContents = append(dirContents, map[string]interface{}{
			"name": name,
			"type": entryType,
			"size": size,
		})
	}

	return resultToJson(FunctionResult{Result: dirContents})
}