fix: better tool output rendering (#562)

Author: jayhack

src/codegen/extensions/langchain/agent.py

@@ -24,7 +24,6 @@
     RevealSymbolTool,
     SearchTool,
     SemanticEditTool,
-    SemanticSearchTool,
     ViewFileTool,
 )
 
@@ -70,8 +69,8 @@ def create_codebase_agent(
         MoveSymbolTool(codebase),
         RevealSymbolTool(codebase),
         SemanticEditTool(codebase),
-        SemanticSearchTool(codebase),
         ReplacementEditTool(codebase),
+        # SemanticSearchTool(codebase),
         # =====[ Github Integration ]=====
         # Enable Github integration
         # GithubCreatePRTool(codebase),

src/codegen/extensions/tools/edit_file.py

@@ -7,7 +7,7 @@
 from codegen.sdk.core.codebase import Codebase
 
 from .observation import Observation
-from .view_file import ViewFileObservation, view_file
+from .replacement_edit import generate_diff
 
 
 class EditFileObservation(Observation):
@@ -16,23 +16,26 @@ class EditFileObservation(Observation):
     filepath: str = Field(
         description="Path to the edited file",
     )
-    file_info: ViewFileObservation = Field(
-        description="Information about the edited file",
+    diff: str = Field(
+        description="Unified diff showing the changes made",
     )
 
     str_template: ClassVar[str] = "Edited file {filepath}"
 
+    def render(self) -> str:
+        """Render edit results in a clean format."""
+        return f"""[EDIT FILE]: {self.filepath}
 
-def edit_file(codebase: Codebase, filepath: str, content: str) -> EditFileObservation:
-    """Edit a file by replacing its entire content.
+{self.diff}"""
+
+
+def edit_file(codebase: Codebase, filepath: str, new_content: str) -> EditFileObservation:
+    """Edit the contents of a file.
 
     Args:
         codebase: The codebase to operate on
-        filepath: Path to the file to edit
-        content: New content for the file
-
-    Returns:
-        EditFileObservation containing updated file state, or error if file not found
+        filepath: Path to the file relative to workspace root
+        new_content: New content for the file
     """
     try:
         file = codebase.get_file(filepath)
@@ -41,52 +44,18 @@ def edit_file(codebase: Codebase, filepath: str, content: str) -> EditFileObserv
             status="error",
             error=f"File not found: {filepath}",
             filepath=filepath,
-            file_info=ViewFileObservation(
-                status="error",
-                error=f"File not found: {filepath}",
-                filepath=filepath,
-                content="",
-                line_count=0,
-            ),
-        )
-
-    if file is None:
-        return EditFileObservation(
-            status="error",
-            error=f"File not found: {filepath}",
-            filepath=filepath,
-            file_info=ViewFileObservation(
-                status="error",
-                error=f"File not found: {filepath}",
-                filepath=filepath,
-                content="",
-                line_count=0,
-            ),
+            diff="",
         )
 
-    try:
-        file.edit(content)
-        codebase.commit()
+    # Generate diff before making changes
+    diff = generate_diff(file.content, new_content)
 
-        # Get updated file info using view_file
-        file_info = view_file(codebase, filepath)
+    # Apply the edit
+    file.edit(new_content)
+    codebase.commit()
 
-        return EditFileObservation(
-            status="success",
-            filepath=filepath,
-            file_info=file_info,
-        )
-
-    except Exception as e:
-        return EditFileObservation(
-            status="error",
-            error=f"Failed to edit file: {e!s}",
-            filepath=filepath,
-            file_info=ViewFileObservation(
-                status="error",
-                error=f"Failed to edit file: {e!s}",
-                filepath=filepath,
-                content="",
-                line_count=0,
-            ),
-        )
+    return EditFileObservation(
+        status="success",
+        filepath=filepath,
+        diff=diff,
+    )

src/codegen/extensions/tools/list_directory.py

@@ -1,81 +1,144 @@
 """Tool for listing directory contents."""
 
-from typing import ClassVar, Union
+from typing import ClassVar
 
-from pydantic import BaseModel, Field
+from pydantic import Field
 
 from codegen.sdk.core.codebase import Codebase
 from codegen.sdk.core.directory import Directory
 
 from .observation import Observation
 
 
-class DirectoryInfo(BaseModel):
+class DirectoryInfo(Observation):
     """Information about a directory."""
 
-    name: str = Field(description="Name of the directory")
-    path: str = Field(description="Full path to the directory")
-    files: list[str] = Field(description="List of files in this directory")
-    subdirectories: list[Union[str, "DirectoryInfo"]] = Field(
-        description="List of subdirectories (either names or full DirectoryInfo objects depending on depth)",
+    name: str = Field(
+        description="Name of the directory",
+    )
+    path: str = Field(
+        description="Full path to the directory",
+    )
+    files: list[str] | None = Field(
+        default=None,
+        description="List of files in this directory (None if at max depth)",
+    )
+    subdirectories: list["DirectoryInfo"] = Field(
+        default_factory=list,
+        description="List of subdirectories",
+    )
+    is_leaf: bool = Field(
+        default=False,
+        description="Whether this is a leaf node (at max depth)",
     )
 
+    str_template: ClassVar[str] = "Directory {path} ({file_count} files, {dir_count} subdirs)"
+
+    def _get_details(self) -> dict[str, int]:
+        """Get details for string representation."""
+        return {
+            "file_count": len(self.files or []),
+            "dir_count": len(self.subdirectories),
+        }
+
+    def render(self) -> str:
+        """Render directory listing as a file tree."""
+        lines = [
+            f"[LIST DIRECTORY]: {self.path}",
+            "",
+        ]
+
+        def add_tree_item(name: str, prefix: str = "", is_last: bool = False) -> tuple[str, str]:
+            """Helper to format a tree item with proper prefix."""
+            marker = "└── " if is_last else "├── "
+            indent = "    " if is_last else "│   "
+            return prefix + marker + name, prefix + indent
+
+        def build_tree(items: list[tuple[str, bool, "DirectoryInfo | None"]], prefix: str = "") -> list[str]:
+            """Recursively build tree with proper indentation."""
+            if not items:
+                return []
+
+            result = []
+            for i, (name, is_dir, dir_info) in enumerate(items):
+                is_last = i == len(items) - 1
+                line, new_prefix = add_tree_item(name, prefix, is_last)
+                result.append(line)
+
+                # If this is a directory and not a leaf node, show its contents
+                if dir_info and not dir_info.is_leaf:
+                    subitems = []
+                    # Add files first
+                    if dir_info.files:
+                        for f in sorted(dir_info.files):
+                            subitems.append((f, False, None))
+                    # Then add subdirectories
+                    for d in dir_info.subdirectories:
+                        subitems.append((d.name + "/", True, d))
+
+                    result.extend(build_tree(subitems, new_prefix))
+
+            return result
+
+        # Sort files and directories
+        items = []
+        if self.files:
+            for f in sorted(self.files):
+                items.append((f, False, None))
+        for d in self.subdirectories:
+            items.append((d.name + "/", True, d))
+
+        if not items:
+            lines.append("(empty directory)")
+            return "\n".join(lines)
+
+        # Generate tree
+        lines.extend(build_tree(items))
+
+        return "\n".join(lines)
+
 
 class ListDirectoryObservation(Observation):
     """Response from listing directory contents."""
 
-    path: str = Field(description="Path to the listed directory")
-    directory_info: DirectoryInfo = Field(description="Information about the directory and its contents")
-    depth: int = Field(description="How deep the directory traversal went")
+    directory_info: DirectoryInfo = Field(
+        description="Information about the directory",
+    )
+
+    str_template: ClassVar[str] = "{directory_info}"
 
-    str_template: ClassVar[str] = "Listed contents of {path} (depth={depth})"
+    def render(self) -> str:
+        """Render directory listing."""
+        return self.directory_info.render()
 
 
-def list_directory(codebase: Codebase, dirpath: str = "./", depth: int = 1) -> ListDirectoryObservation:
+def list_directory(codebase: Codebase, path: str = "./", depth: int = 2) -> ListDirectoryObservation:
     """List contents of a directory.
 
     Args:
         codebase: The codebase to operate on
-        dirpath: Path to directory relative to workspace root
+        path: Path to directory relative to workspace root
         depth: How deep to traverse the directory tree. Default is 1 (immediate children only).
                Use -1 for unlimited depth.
-
-    Returns:
-        ListDirectoryObservation containing directory contents and metadata
     """
     try:
-        directory = codebase.get_directory(dirpath)
+        directory = codebase.get_directory(path)
     except ValueError:
         return ListDirectoryObservation(
             status="error",
-            error=f"Directory not found: {dirpath}",
-            path=dirpath,
+            error=f"Directory not found: {path}",
             directory_info=DirectoryInfo(
-                name="",
-                path=dirpath,
+                status="error",
+                name=path.split("/")[-1],
+                path=path,
                 files=[],
                 subdirectories=[],
             ),
-            depth=depth,
-        )
-
-    if not directory:
-        return ListDirectoryObservation(
-            status="error",
-            error=f"Directory not found: {dirpath}",
-            path=dirpath,
-            directory_info=DirectoryInfo(
-                name="",
-                path=dirpath,
-                files=[],
-                subdirectories=[],
-            ),
-            depth=depth,
         )
 
     def get_directory_info(dir_obj: Directory, current_depth: int) -> DirectoryInfo:
         """Helper function to get directory info recursively."""
-        # Get direct files
+        # Get direct files (always include files unless at max depth)
         all_files = []
         for file in dir_obj.files:
             if file.directory == dir_obj:
@@ -86,38 +149,32 @@ def get_directory_info(dir_obj: Directory, current_depth: int) -> DirectoryInfo:
         for subdir in dir_obj.subdirectories:
             # Only include direct descendants
             if subdir.parent == dir_obj:
-                if current_depth != 1:
+                if current_depth > 1 or current_depth == -1:
+                    # For deeper traversal, get full directory info
                     new_depth = current_depth - 1 if current_depth > 1 else -1
                     subdirs.append(get_directory_info(subdir, new_depth))
                 else:
-                    # At max depth, just include name
-                    subdirs.append(subdir.name)
+                    # At max depth, return a leaf node
+                    subdirs.append(
+                        DirectoryInfo(
+                            status="success",
+                            name=subdir.name,
+                            path=subdir.dirpath,
+                            files=None,  # Don't include files at max depth
+                            is_leaf=True,
+                        )
+                    )
 
         return DirectoryInfo(
+            status="success",
             name=dir_obj.name,
             path=dir_obj.dirpath,
-            files=all_files,
+            files=sorted(all_files),
             subdirectories=subdirs,
         )
 
-    try:
-        directory_info = get_directory_info(directory, depth)
-        return ListDirectoryObservation(
-            status="success",
-            path=dirpath,
-            directory_info=directory_info,
-            depth=depth,
-        )
-    except Exception as e:
-        return ListDirectoryObservation(
-            status="error",
-            error=f"Failed to list directory: {e!s}",
-            path=dirpath,
-            directory_info=DirectoryInfo(
-                name="",
-                path=dirpath,
-                files=[],
-                subdirectories=[],
-            ),
-            depth=depth,
-        )
+    dir_info = get_directory_info(directory, depth)
+    return ListDirectoryObservation(
+        status="success",
+        directory_info=dir_info,
+    )