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

View File

@ -39,7 +39,7 @@ var DirTreeTool = model.Tool{
{ {
Name: "depth", Name: "depth",
Type: "integer", 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) { 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 { if depth, err = strconv.Atoi(v); err != nil {
return "", fmt.Errorf("invalid `depth` value, expected integer but got string that cannot convert: %v", tmp) return "", fmt.Errorf("invalid `depth` value, expected integer but got string that cannot convert: %v", tmp)
} }
depth = max(0, min(5, depth))
default: default:
return "", fmt.Errorf("expected int or string for max_depth, got %T", tmp) 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()) output.WriteString(prefix + branch + file.Name())
if file.IsDir() { if file.IsDir() {
output.WriteString("/\n") output.WriteString("/\n")
if depth != 0 { if depth > 0 {
var nextPrefix string var nextPrefix string
if isLast { if isLast {
nextPrefix = prefix + " " nextPrefix = prefix + " "