ml-knowledge-platform/knowledge_platform/tools/builtin/tree.py
2026-02-16 04:50:51 -08:00

132 lines
3.9 KiB
Python

"""TreeTool — visualize directory structure as an indented tree.
Produces a hierarchical view of a directory, respecting depth limits
and hidden-file filtering. Essential for project orientation.
"""
from __future__ import annotations
from pathlib import Path
from typing import Any, ClassVar
from ..base import Tool, ToolParameter, ToolResult
class TreeTool(Tool):
"""Visualize directory structure as an indented tree.
Produces output similar to the ``tree`` command with configurable
depth limit and hidden file filtering.
"""
name: ClassVar[str] = "tree"
description: ClassVar[str] = (
"Display directory structure as an indented tree. "
"Supports depth limiting and hidden file filtering."
)
parameters: ClassVar[list[ToolParameter]] = [
ToolParameter(
name="path",
type="string",
description="Absolute path to the root directory",
),
ToolParameter(
name="max_depth",
type="integer",
description="Maximum depth to recurse (default 3)",
required=False,
default=3,
),
ToolParameter(
name="show_hidden",
type="boolean",
description="Include hidden files/directories",
required=False,
default=False,
),
]
async def execute(self, **kwargs: Any) -> ToolResult:
root = Path(kwargs["path"])
max_depth: int = kwargs.get("max_depth", 3)
show_hidden: bool = kwargs.get("show_hidden", False)
if not root.is_absolute():
return ToolResult.fail(f"Path must be absolute: {root}")
if not root.exists():
return ToolResult.fail(f"Directory not found: {root}")
if not root.is_dir():
return ToolResult.fail(f"Not a directory: {root}")
lines: list[str] = [f"{root.name}/"]
dir_count = 0
file_count = 0
dir_count, file_count = self._walk(
root, lines, "", max_depth, 0, show_hidden
)
output = "\n".join(lines)
summary = f"\n\n{dir_count} directories, {file_count} files"
output += summary
return ToolResult.success(
output,
directories=dir_count,
files=file_count,
)
def _walk(
self,
directory: Path,
lines: list[str],
prefix: str,
max_depth: int,
current_depth: int,
show_hidden: bool,
) -> tuple[int, int]:
"""Recursively build tree lines.
Returns (directory_count, file_count).
"""
if current_depth >= max_depth:
return 0, 0
try:
entries = sorted(directory.iterdir(), key=lambda p: (p.is_file(), p.name))
except PermissionError:
lines.append(f"{prefix}[permission denied]")
return 0, 0
if not show_hidden:
entries = [e for e in entries if not e.name.startswith(".")]
dir_count = 0
file_count = 0
total = len(entries)
for i, entry in enumerate(entries):
is_last = i == total - 1
connector = "└── " if is_last else "├── "
extension = " " if is_last else ""
if entry.is_dir():
lines.append(f"{prefix}{connector}{entry.name}/")
dir_count += 1
sub_dirs, sub_files = self._walk(
entry,
lines,
prefix + extension,
max_depth,
current_depth + 1,
show_hidden,
)
dir_count += sub_dirs
file_count += sub_files
elif entry.is_file():
lines.append(f"{prefix}{connector}{entry.name}")
file_count += 1
return dir_count, file_count