144 lines
3.1 KiB
Go
144 lines
3.1 KiB
Go
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
|
|
}
|
|
|