Add dir_tree tool

This commit is contained in:
Matt Low 2024-03-22 20:30:34 +00:00
parent 91c74d9e1e
commit c51644e78e
2 changed files with 144 additions and 0 deletions

143
pkg/lmcli/tools/dir_tree.go Normal file
View File

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

View File

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