diff --git a/pkg/lmcli/tools/dir_tree.go b/pkg/lmcli/tools/dir_tree.go new file mode 100644 index 0000000..5642629 --- /dev/null +++ b/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 +} + diff --git a/pkg/lmcli/tools/tools.go b/pkg/lmcli/tools/tools.go index 07dfa82..87f5b8f 100644 --- a/pkg/lmcli/tools/tools.go +++ b/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,