lmcli/pkg/agent/toolbox/dir_tree.go

package toolbox

import (
	"fmt"
	"os"
	"path/filepath"
	"strconv"
	"strings"

	toolutil "git.mlow.ca/mlow/lmcli/pkg/agent/toolbox/util"
	"git.mlow.ca/mlow/lmcli/pkg/api"
)

const TREE_DESCRIPTION = `Retrieve a tree-like view of a directory's contents.

Use these results for your own reference in completing your task, they do not need to be shown to the user.

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 = api.ToolSpec{
	Name:        "dir_tree",
	Description: TREE_DESCRIPTION,
	Parameters: []api.ToolParameter{
		{
			Name:        "relative_path",
			Type:        "string",
			Description: "If set, display the tree starting from this path relative to the current one.",
		},
		{
			Name:        "depth",
			Type:        "integer",
			Description: "Depth of directory recursion. Defaults to 0 (no recursion), maximum of 5.",
		},
	},
	Impl: func(tool *api.ToolSpec, args map[string]interface{}) (string, error) {
		var relativeDir string
		if tmp, ok := args["relative_path"]; ok {
			relativeDir, ok = tmp.(string)
			if !ok {
				return "", fmt.Errorf("expected string for relative_path, got %T", tmp)
			}
		}

		var depth int = 0 // Default value if not provided
		if tmp, ok := args["depth"]; ok {
			switch v := tmp.(type) {
			case float64:
				depth = int(v)
			case string:
				var err error
				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)
			}
		}

		result := tree(relativeDir, depth)
		ret, err := result.ToJson()
		if err != nil {
			return "", fmt.Errorf("could not serialize result: %v", err)
		}
		return ret, nil
	},
}

func tree(path string, depth int) api.CallResult {
	if path == "" {
		path = "."
	}
	ok, reason := toolutil.IsPathWithinCWD(path)
	if !ok {
		return api.CallResult{Message: reason}
	}

	var treeOutput strings.Builder
	treeOutput.WriteString(path + "\n")
	err := buildTree(&treeOutput, path, "", depth)
	if err != nil {
		return api.CallResult{
			Message: err.Error(),
		}
	}

	return api.CallResult{Result: treeOutput.String()}
}

func buildTree(output *strings.Builder, path string, prefix string, depth 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 depth > 0 {
				var nextPrefix string
				if isLast {
					nextPrefix = prefix + "    "
				} else {
					nextPrefix = prefix + "│   "
				}
				buildTree(output, filepath.Join(path, file.Name()), nextPrefix, depth-1)
			}
		} else {
			output.WriteString(sizeStr + "\n")
		}
	}

	return nil
}