Add dir_tree tool

pkg/lmcli/tools/dir_tree.go

@@ -0,0 +1,143 @@
+package tools
+
+import (
+	"fmt"
+	"os"
+	"path/filepath"
+	"strconv"
+	"strings"
+
+	"git.mlow.ca/mlow/lmcli/pkg/lmcli/model"
+	toolutil "git.mlow.ca/mlow/lmcli/pkg/lmcli/tools/util"
+)
+
+const TREE_DESCRIPTION = `Retrieve a tree view of a directory's contents.
+
+Example result:
+{
+	"message": "success",
+	"result": ".
+├── a_directory/
+│   ├── file1.txt (100 bytes)
+│   └── file2.txt (200 bytes)
+├── a_file.txt (123 bytes)
+└── another_file.txt (456 bytes)"
+}
+`
+
+var DirTreeTool = model.Tool{
+	Name:        "dir_tree",
+	Description: TREE_DESCRIPTION,
+	Parameters: []model.ToolParameter{
+		{
+			Name:        "relative_path",
+			Type:        "string",
+			Description: "If set, display the tree starting from this path relative to the current one.",
+		},
+		{
+			Name:        "max_depth",
+			Type:        "integer",
+			Description: "Maximum depth of recursion. Default is unlimited.",
+		},
+	},
+	Impl: func(tool *model.Tool, 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)
+			}
+		}
+
+		var maxDepth int = -1
+		tmp, ok = args["max_depth"]
+		if ok {
+			maxDepth, ok = tmp.(int)
+			if !ok {
+				if tmps, ok := tmp.(string); ok {
+					tmpi, err := strconv.Atoi(tmps)
+					maxDepth = tmpi
+					if err != nil  {
+						return "", fmt.Errorf("Invalid max_depth in function arguments: %v", tmp)
+					}
+				} else {
+					return "", fmt.Errorf("Invalid max_depth in function arguments: %v", tmp)
+				}
+			}
+		}
+
+		result := tree(relativeDir, maxDepth)
+		ret, err := result.ToJson()
+		if err != nil {
+			return "", fmt.Errorf("Could not serialize result: %v", err)
+		}
+		return ret, nil
+	},
+}
+
+func tree(path string, maxDepth int) model.CallResult {
+	if path == "" {
+		path = "."
+	}
+	ok, reason := toolutil.IsPathWithinCWD(path)
+	if !ok {
+		return model.CallResult{Message: reason}
+	}
+
+	var treeOutput strings.Builder
+	treeOutput.WriteString(path + "\n")
+	err := buildTree(&treeOutput, path, "", maxDepth)
+	if err != nil {
+		return model.CallResult{
+			Message: err.Error(),
+		}
+	}
+
+	return model.CallResult{Result: treeOutput.String()}
+}
+
+func buildTree(output *strings.Builder, path string, prefix string, maxDepth int) error {
+	files, err := os.ReadDir(path)
+	if err != nil {
+		return err
+	}
+
+	for i, file := range files {
+		if strings.HasPrefix(file.Name(), ".") {
+			// Skip hidden files and directories
+			continue
+		}
+
+		isLast := i == len(files)-1
+		var branch string
+		if isLast {
+			branch = "└── "
+		} else {
+			branch = "├── "
+		}
+
+		info, _ := file.Info()
+		size := info.Size()
+		sizeStr := fmt.Sprintf(" (%d bytes)", size)
+
+		output.WriteString(prefix + branch + file.Name())
+		if file.IsDir() {
+			output.WriteString("/\n")
+			if maxDepth != 0 {
+				var nextPrefix string
+				if isLast {
+					nextPrefix = prefix + "    "
+				} else {
+					nextPrefix = prefix + "│   "
+				}
+				buildTree(output, filepath.Join(path, file.Name()), nextPrefix, maxDepth-1)
+			}
+		} else {
+			output.WriteString(sizeStr + "\n")
+		}
+	}
+
+	return nil
+}
+

pkg/lmcli/tools/tools.go

@@ -7,6 +7,7 @@ import (
 )
 
 var AvailableTools map[string]model.Tool = map[string]model.Tool{
+	"dir_tree":           DirTreeTool,
 	"read_dir":           ReadDirTool,
 	"read_file":          ReadFileTool,
 	"write_file":         WriteFileTool,