Update dir_tree to have maximum depth of 5

Until we have some mechanism in place for confirming tool calls with the
user before executing, it's dangerous to allow unlimited depth
This commit is contained in:
Matt Low 2024-05-21 00:08:17 +00:00
parent 8c53752146
commit 5c1248184b
1 changed files with 3 additions and 2 deletions

View File

@ -39,7 +39,7 @@ var DirTreeTool = model.Tool{
{
Name: "depth",
Type: "integer",
Description: "Depth of directory recursion. Default 0. Use -1 for unlimited.",
Description: "Depth of directory recursion. Defaults to 0 (no recursion), maximum of 5.",
},
},
Impl: func(tool *model.Tool, args map[string]interface{}) (string, error) {
@ -61,6 +61,7 @@ var DirTreeTool = model.Tool{
if depth, err = strconv.Atoi(v); err != nil {
return "", fmt.Errorf("invalid `depth` value, expected integer but got string that cannot convert: %v", tmp)
}
depth = max(0, min(5, depth))
default:
return "", fmt.Errorf("expected int or string for max_depth, got %T", tmp)
}
@ -123,7 +124,7 @@ func buildTree(output *strings.Builder, path string, prefix string, depth int) e
output.WriteString(prefix + branch + file.Name())
if file.IsDir() {
output.WriteString("/\n")
if depth != 0 {
if depth > 0 {
var nextPrefix string
if isLast {
nextPrefix = prefix + " "