Apply changes from commit 046b238

Original commit by Tawsif Kamal: Revert "Revert "Adding Schema for Tool Outputs"" (codegen-sh#894)

Reverts codegen-sh#892

---------

Co-authored-by: Rushil Patel <[email protected]>
Co-authored-by: rushilpatel0 <[email protected]>

Author: Zeeeepa

src/codegen/agents/data.py

@@ -54,6 +54,7 @@ class ToolMessageData(BaseMessage):
     tool_id: Optional[str] = None
     status: Optional[str] = None
 
+
 @dataclass
 class FunctionMessageData(BaseMessage):
     """Represents a function message."""

src/codegen/agents/tracer.py

@@ -78,7 +78,8 @@ def extract_structured_data(self, chunk: dict[str, Any]) -> Optional[BaseMessage
                 tool_response=getattr(latest_message, "artifact", content),
                 tool_id=getattr(latest_message, "tool_call_id", None),
                 status=getattr(latest_message, "status", None),
-            )        elif message_type == "function":
+            )
+        elif message_type == "function":
             return FunctionMessageData(type=message_type, content=content)
         else:
             return UnknownMessage(type=message_type, content=content)

src/codegen/extensions/langchain/tools.py

@@ -1,8 +1,10 @@
 """Langchain tools for workspace operations."""
 
 from collections.abc import Callable
-from typing import ClassVar, Literal
+from typing import Annotated, ClassVar, Literal, Optional
 
+from langchain_core.messages import ToolMessage
+from langchain_core.tools import InjectedToolCallId
 from langchain_core.tools.base import BaseTool
 from pydantic import BaseModel, Field
 
@@ -52,10 +54,11 @@ class ViewFileInput(BaseModel):
     """Input for viewing a file."""
 
     filepath: str = Field(..., description="Path to the file relative to workspace root")
-    start_line: int | None = Field(None, description="Starting line number to view (1-indexed, inclusive)")
-    end_line: int | None = Field(None, description="Ending line number to view (1-indexed, inclusive)")
-    max_lines: int | None = Field(None, description="Maximum number of lines to view at once, defaults to 500")
-    line_numbers: bool | None = Field(True, description="If True, add line numbers to the content (1-indexed)")
+    start_line: Optional[int] = Field(None, description="Starting line number to view (1-indexed, inclusive)")
+    end_line: Optional[int] = Field(None, description="Ending line number to view (1-indexed, inclusive)")
+    max_lines: Optional[int] = Field(None, description="Maximum number of lines to view at once, defaults to 500")
+    line_numbers: Optional[bool] = Field(True, description="If True, add line numbers to the content (1-indexed)")
+    tool_call_id: Annotated[str, InjectedToolCallId]
 
 
 class ViewFileTool(BaseTool):
@@ -73,12 +76,13 @@ def __init__(self, codebase: Codebase) -> None:
 
     def _run(
         self,
+        tool_call_id: str,
         filepath: str,
-        start_line: int | None = None,
-        end_line: int | None = None,
-        max_lines: int | None = None,
-        line_numbers: bool | None = True,
-    ) -> str:
+        start_line: Optional[int] = None,
+        end_line: Optional[int] = None,
+        max_lines: Optional[int] = None,
+        line_numbers: Optional[bool] = True,
+    ) -> ToolMessage:
         result = view_file(
             self.codebase,
             filepath,
@@ -88,14 +92,15 @@ def _run(
             max_lines=max_lines if max_lines is not None else 500,
         )
 
-        return result.render()
+        return result.render(tool_call_id)
 
 
 class ListDirectoryInput(BaseModel):
     """Input for listing directory contents."""
 
     dirpath: str = Field(default="./", description="Path to directory relative to workspace root")
     depth: int = Field(default=1, description="How deep to traverse. Use -1 for unlimited depth.")
+    tool_call_id: Annotated[str, InjectedToolCallId]
 
 
 class ListDirectoryTool(BaseTool):
@@ -109,9 +114,9 @@ class ListDirectoryTool(BaseTool):
     def __init__(self, codebase: Codebase) -> None:
         super().__init__(codebase=codebase)
 
-    def _run(self, dirpath: str = "./", depth: int = 1) -> str:
+    def _run(self, tool_call_id: str, dirpath: str = "./", depth: int = 1) -> ToolMessage:
         result = list_directory(self.codebase, dirpath, depth)
-        return result.render()
+        return result.render(tool_call_id)
 
 
 class SearchInput(BaseModel):
@@ -126,6 +131,7 @@ class SearchInput(BaseModel):
     page: int = Field(default=1, description="Page number to return (1-based, default: 1)")
     files_per_page: int = Field(default=10, description="Number of files to return per page (default: 10)")
     use_regex: bool = Field(default=False, description="Whether to treat query as a regex pattern (default: False)")
+    tool_call_id: Annotated[str, InjectedToolCallId]
 
 
 class SearchTool(BaseTool):
@@ -139,16 +145,17 @@ class SearchTool(BaseTool):
     def __init__(self, codebase: Codebase) -> None:
         super().__init__(codebase=codebase)
 
-    def _run(self, query: str, file_extensions: list[str] | None = None, page: int = 1, files_per_page: int = 10, use_regex: bool = False) -> str:
+    def _run(self, tool_call_id: str, query: str, file_extensions: Optional[list[str]] = None, page: int = 1, files_per_page: int = 10, use_regex: bool = False) -> ToolMessage:
         result = search(self.codebase, query, file_extensions=file_extensions, page=page, files_per_page=files_per_page, use_regex=use_regex)
-        return result.render()
+        return result.render(tool_call_id)
 
 
 class EditFileInput(BaseModel):
     """Input for editing a file."""
 
     filepath: str = Field(..., description="Path to the file to edit")
     content: str = Field(..., description="New content for the file")
+    tool_call_id: Annotated[str, InjectedToolCallId]
 
 
 class EditFileTool(BaseTool):
@@ -181,9 +188,9 @@ class EditFileTool(BaseTool):
     def __init__(self, codebase: Codebase) -> None:
         super().__init__(codebase=codebase)
 
-    def _run(self, filepath: str, content: str) -> str:
+    def _run(self, filepath: str, content: str, tool_call_id: str) -> str:
         result = edit_file(self.codebase, filepath, content)
-        return result.render()
+        return result.render(tool_call_id)
 
 
 class CreateFileInput(BaseModel):
@@ -340,6 +347,7 @@ class SemanticEditInput(BaseModel):
     edit_content: str = Field(..., description=FILE_EDIT_PROMPT)
     start: int = Field(default=1, description="Starting line number (1-indexed, inclusive). Default is 1.")
     end: int = Field(default=-1, description="Ending line number (1-indexed, inclusive). Default is -1 (end of file).")
+    tool_call_id: Annotated[str, InjectedToolCallId]
 
 
 class SemanticEditTool(BaseTool):
@@ -353,10 +361,10 @@ class SemanticEditTool(BaseTool):
     def __init__(self, codebase: Codebase) -> None:
         super().__init__(codebase=codebase)
 
-    def _run(self, filepath: str, edit_content: str, start: int = 1, end: int = -1) -> str:
+    def _run(self, filepath: str, tool_call_id: str, edit_content: str, start: int = 1, end: int = -1) -> ToolMessage:
         # Create the the draft editor mini llm
         result = semantic_edit(self.codebase, filepath, edit_content, start=start, end=end)
-        return result.render()
+        return result.render(tool_call_id)
 
 
 class RenameFileInput(BaseModel):
@@ -1033,6 +1041,7 @@ class RelaceEditInput(BaseModel):
 
     filepath: str = Field(..., description="Path of the file relative to workspace root")
     edit_snippet: str = Field(..., description=RELACE_EDIT_PROMPT)
+    tool_call_id: Annotated[str, InjectedToolCallId]
 
 
 class RelaceEditTool(BaseTool):
@@ -1046,9 +1055,9 @@ class RelaceEditTool(BaseTool):
     def __init__(self, codebase: Codebase) -> None:
         super().__init__(codebase=codebase)
 
-    def _run(self, filepath: str, edit_snippet: str) -> str:
+    def _run(self, filepath: str, edit_snippet: str, tool_call_id: str) -> ToolMessage:
         result = relace_edit(self.codebase, filepath, edit_snippet)
-        return result.render()
+        return result.render(tool_call_id=tool_call_id)
 
 
 class ReflectionInput(BaseModel):

src/codegen/extensions/tools/edit_file.py

@@ -1,32 +1,53 @@
 """Tool for editing file contents."""
 
-from typing import ClassVar
+from typing import TYPE_CHECKING, ClassVar, Optional
 
+from langchain_core.messages import ToolMessage
 from pydantic import Field
 
 from codegen.sdk.core.codebase import Codebase
 
 from .observation import Observation
 from .replacement_edit import generate_diff
 
+if TYPE_CHECKING:
+    from .tool_output_types import EditFileArtifacts
+
 
 class EditFileObservation(Observation):
     """Response from editing a file."""
 
     filepath: str = Field(
         description="Path to the edited file",
     )
-    diff: str = Field(
+    diff: Optional[str] = Field(
+        default=None,
         description="Unified diff showing the changes made",
     )
 
     str_template: ClassVar[str] = "Edited file {filepath}"
 
-    def render(self) -> str:
+    def render(self, tool_call_id: str) -> ToolMessage:
         """Render edit results in a clean format."""
-        return f"""[EDIT FILE]: {self.filepath}
-
-{self.diff}"""
+        if self.status == "error":
+            artifacts_error: EditFileArtifacts = {"filepath": self.filepath, "error": self.error}
+            return ToolMessage(
+                content=f"[ERROR EDITING FILE]: {self.filepath}: {self.error}",
+                status=self.status,
+                name="edit_file",
+                artifact=artifacts_error,
+                tool_call_id=tool_call_id,
+            )
+
+        artifacts_success: EditFileArtifacts = {"filepath": self.filepath, "diff": self.diff}
+
+        return ToolMessage(
+            content=f"""[EDIT FILE]: {self.filepath}\n\n{self.diff}""",
+            status=self.status,
+            name="edit_file",
+            artifact=artifacts_success,
+            tool_call_id=tool_call_id,
+        )
 
 
 def edit_file(codebase: Codebase, filepath: str, new_content: str) -> EditFileObservation:

src/codegen/extensions/tools/list_directory.py

@@ -2,13 +2,14 @@
 
 from typing import ClassVar
 
+from langchain_core.messages import ToolMessage
 from pydantic import Field
 
+from codegen.extensions.tools.observation import Observation
+from codegen.extensions.tools.tool_output_types import ListDirectoryArtifacts
 from codegen.sdk.core.codebase import Codebase
 from codegen.sdk.core.directory import Directory
 
-from .observation import Observation
-
 
 class DirectoryInfo(Observation):
     """Information about a directory."""
@@ -31,6 +32,14 @@ class DirectoryInfo(Observation):
         default=False,
         description="Whether this is a leaf node (at max depth)",
     )
+    depth: int = Field(
+        default=0,
+        description="Current depth in the tree",
+    )
+    max_depth: int = Field(
+        default=1,
+        description="Maximum depth allowed",
+    )
 
     str_template: ClassVar[str] = "Directory {path} ({file_count} files, {dir_count} subdirs)"
 
@@ -41,7 +50,7 @@ def _get_details(self) -> dict[str, int]:
             "dir_count": len(self.subdirectories),
         }
 
-    def render(self) -> str:
+    def render_as_string(self) -> str:
         """Render directory listing as a file tree."""
         lines = [
             f"[LIST DIRECTORY]: {self.path}",
@@ -97,6 +106,26 @@ def build_tree(items: list[tuple[str, bool, "DirectoryInfo | None"]], prefix: st
 
         return "\n".join(lines)
 
+    def to_artifacts(self) -> ListDirectoryArtifacts:
+        """Convert directory info to artifacts for UI."""
+        artifacts: ListDirectoryArtifacts = {
+            "dirpath": self.path,
+            "name": self.name,
+            "is_leaf": self.is_leaf,
+            "depth": self.depth,
+            "max_depth": self.max_depth,
+        }
+
+        if self.files is not None:
+            artifacts["files"] = self.files
+            artifacts["file_paths"] = [f"{self.path}/{f}" for f in self.files]
+
+        if self.subdirectories:
+            artifacts["subdirs"] = [d.name for d in self.subdirectories]
+            artifacts["subdir_paths"] = [d.path for d in self.subdirectories]
+
+        return artifacts
+
 
 class ListDirectoryObservation(Observation):
     """Response from listing directory contents."""
@@ -107,9 +136,29 @@ class ListDirectoryObservation(Observation):
 
     str_template: ClassVar[str] = "{directory_info}"
 
-    def render(self) -> str:
-        """Render directory listing."""
-        return self.directory_info.render()
+    def render(self, tool_call_id: str) -> ToolMessage:
+        """Render directory listing with artifacts for UI."""
+        if self.status == "error":
+            error_artifacts: ListDirectoryArtifacts = {
+                "dirpath": self.directory_info.path,
+                "name": self.directory_info.name,
+                "error": self.error,
+            }
+            return ToolMessage(
+                content=f"[ERROR LISTING DIRECTORY]: {self.directory_info.path}: {self.error}",
+                status=self.status,
+                name="list_directory",
+                artifact=error_artifacts,
+                tool_call_id=tool_call_id,
+            )
+
+        return ToolMessage(
+            content=self.directory_info.render_as_string(),
+            status=self.status,
+            name="list_directory",
+            artifact=self.directory_info.to_artifacts(),
+            tool_call_id=tool_call_id,
+        )
 
 
 def list_directory(codebase: Codebase, path: str = "./", depth: int = 2) -> ListDirectoryObservation:
@@ -136,7 +185,7 @@ def list_directory(codebase: Codebase, path: str = "./", depth: int = 2) -> List
             ),
         )
 
-    def get_directory_info(dir_obj: Directory, current_depth: int) -> DirectoryInfo:
+    def get_directory_info(dir_obj: Directory, current_depth: int, max_depth: int) -> DirectoryInfo:
         """Helper function to get directory info recursively."""
         # Get direct files (always include files unless at max depth)
         all_files = []
@@ -151,7 +200,7 @@ def get_directory_info(dir_obj: Directory, current_depth: int) -> DirectoryInfo:
                 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))
+                    subdirs.append(get_directory_info(subdir, new_depth, max_depth))
                 else:
                     # At max depth, return a leaf node
                     subdirs.append(
@@ -161,6 +210,8 @@ def get_directory_info(dir_obj: Directory, current_depth: int) -> DirectoryInfo:
                             path=subdir.dirpath,
                             files=None,  # Don't include files at max depth
                             is_leaf=True,
+                            depth=current_depth,
+                            max_depth=max_depth,
                         )
                     )
 
@@ -170,9 +221,11 @@ def get_directory_info(dir_obj: Directory, current_depth: int) -> DirectoryInfo:
             path=dir_obj.dirpath,
             files=sorted(all_files),
             subdirectories=subdirs,
+            depth=current_depth,
+            max_depth=max_depth,
         )
 
-    dir_info = get_directory_info(directory, depth)
+    dir_info = get_directory_info(directory, depth, depth)
     return ListDirectoryObservation(
         status="success",
         directory_info=dir_info,