Merge pull request #70 from nforro/tools

nforro · web-flow · commit c66c300dd510 · 2025-08-06T09:13:57.000+02:00
Add basic text editor tools
diff --git a/beeai/agents/tests/unit/test_tools.py b/beeai/agents/tests/unit/test_tools.py
@@ -16,6 +16,16 @@
     SetZStreamReleaseTool,
     SetZStreamReleaseToolInput,
 )
+from tools.text import (
+    CreateTool,
+    CreateToolInput,
+    ViewTool,
+    ViewToolInput,
+    InsertTool,
+    InsertToolInput,
+    StrReplaceTool,
+    StrReplaceToolInput,
+)
 
 
 @pytest.mark.parametrize(
@@ -149,3 +159,123 @@ async def test_set_zstream_release(autorelease_spec):
     result = output.result
     assert not result
     assert autorelease_spec.read_text().splitlines()[3] == "Release:        4%{?dist}.%{autorelease -n}"
+
+
+@pytest.mark.asyncio
+async def test_create(tmp_path):
+    test_file = tmp_path / "test.txt"
+    content = "Line 1\nLine 2\n"
+    tool = CreateTool()
+    output = await tool.run(input=CreateToolInput(file=test_file, content=content)).middleware(
+        GlobalTrajectoryMiddleware(pretty=True)
+    )
+    result = output.result
+    assert not result
+    assert test_file.read_text() == content
+
+
+@pytest.mark.asyncio
+async def test_view(tmp_path):
+    test_dir = tmp_path
+    test_file = test_dir / "test.txt"
+    content = "Line 1\nLine 2\nLine 3\n"
+    test_file.write_text(content)
+    tool = ViewTool()
+    output = await tool.run(input=ViewToolInput(path=test_dir)).middleware(
+        GlobalTrajectoryMiddleware(pretty=True)
+    )
+    result = output.result
+    assert result == "test.txt\n"
+    output = await tool.run(input=ViewToolInput(path=test_file)).middleware(
+        GlobalTrajectoryMiddleware(pretty=True)
+    )
+    result = output.result
+    assert result == content
+    output = await tool.run(input=ViewToolInput(path=test_file, view_range=(2, -1))).middleware(
+        GlobalTrajectoryMiddleware(pretty=True)
+    )
+    result = output.result
+    assert (
+        result
+        == dedent(
+            """
+            Line 2
+            Line 3
+            """
+        )[1:]
+    )
+    output = await tool.run(input=ViewToolInput(path=test_file, view_range=(1, 2))).middleware(
+        GlobalTrajectoryMiddleware(pretty=True)
+    )
+    result = output.result
+    assert (
+        result
+        == dedent(
+            """
+            Line 1
+            Line 2
+            """
+        )[1:]
+    )
+
+
+@pytest.mark.parametrize(
+    "line, content",
+    [
+        (
+            0,
+            dedent(
+                """
+                Inserted line
+                Line 1
+                Line 2
+                Line 3
+                """
+            )[1:],
+        ),
+        (
+            1,
+            dedent(
+                """
+                Line 1
+                Inserted line
+                Line 2
+                Line 3
+                """
+            )[1:],
+        ),
+    ],
+)
+@pytest.mark.asyncio
+async def test_insert(line, content, tmp_path):
+    test_file = tmp_path / "test.txt"
+    test_file.write_text("Line 1\nLine 2\nLine 3\n")
+    tool = InsertTool()
+    output = await tool.run(
+        input=InsertToolInput(file=test_file, line=line, new_string="Inserted line")
+    ).middleware(GlobalTrajectoryMiddleware(pretty=True))
+    result = output.result
+    assert not result
+    assert test_file.read_text() == content
+
+
+@pytest.mark.asyncio
+async def test_str_replace(tmp_path):
+    test_file = tmp_path / "test.txt"
+    test_file.write_text("Line 1\nLine 2\nLine 3\n")
+    tool = StrReplaceTool()
+    output = await tool.run(
+        input=StrReplaceToolInput(file=test_file, old_string="Line 2", new_string="LINE_2")
+    ).middleware(GlobalTrajectoryMiddleware(pretty=True))
+    result = output.result
+    assert not result
+    assert (
+        test_file.read_text()
+        == dedent(
+            """
+            Line 1
+            LINE_2
+            Line 3
+            """
+        )[1:]
+    )
diff --git a/beeai/agents/tools/text.py b/beeai/agents/tools/text.py
@@ -0,0 +1,146 @@
+import asyncio
+from pathlib import Path
+
+from pydantic import BaseModel, Field
+
+from beeai_framework.context import RunContext
+from beeai_framework.emitter import Emitter
+from beeai_framework.tools import StringToolOutput, Tool, ToolRunOptions
+
+
+class CreateToolInput(BaseModel):
+    file: Path = Field(description="Absolute path to a file to create")
+    content: str = Field(description="Content to write to the new file")
+
+
+class CreateTool(Tool[CreateToolInput, ToolRunOptions, StringToolOutput]):
+    name = "create"
+    description = """
+    Creates a new file with the specified content.
+    Returns error message on failure.
+    """
+    input_schema = CreateToolInput
+
+    def _create_emitter(self) -> Emitter:
+        return Emitter.root().child(
+            namespace=["tool", "text", self.name],
+            creator=self,
+        )
+
+    async def _run(
+        self, tool_input: CreateToolInput, options: ToolRunOptions | None, context: RunContext
+    ) -> StringToolOutput:
+        try:
+            await asyncio.to_thread(tool_input.file.write_text, tool_input.content)
+        except Exception as e:
+            return StringToolOutput(result=f"Failed to create file: {e}")
+        return StringToolOutput()
+
+
+class ViewToolInput(BaseModel):
+    path: Path = Field(description="Absolute path to a file or directory to view")
+    view_range: tuple[int, int] | None = Field(
+        description="""
+        Tuple of two integers specifying the start and end line numbers to view.
+        Line numbers are 1-indexed, and -1 for the end line means read to the end of the file.
+        This argument only applies when viewing files, not directories.
+        """,
+        default=None,
+    )
+
+
+class ViewTool(Tool[ViewToolInput, ToolRunOptions, StringToolOutput]):
+    name = "view"
+    description = """
+    Outputs the contents of a file or lists the contents of a directory. Can read an entire file
+    or a specific range of lines. Returns error message on failure.
+    """
+    input_schema = ViewToolInput
+
+    def _create_emitter(self) -> Emitter:
+        return Emitter.root().child(
+            namespace=["tool", "text", self.name],
+            creator=self,
+        )
+
+    async def _run(
+        self, tool_input: ViewToolInput, options: ToolRunOptions | None, context: RunContext
+    ) -> StringToolOutput:
+        try:
+            if tool_input.path.is_file():
+                content = await asyncio.to_thread(tool_input.path.read_text)
+                if tool_input.view_range is not None:
+                    start, end = tool_input.view_range
+                    lines = content.splitlines(keepends=True)
+                    content = "".join(lines[start - 1 : None if end < 0 else end])
+                return StringToolOutput(result=content)
+            return StringToolOutput(result="\n".join(e.name for e in tool_input.path.iterdir()) + "\n")
+        except Exception as e:
+            return StringToolOutput(result=f"Failed to view path: {e}")
+
+
+class InsertToolInput(BaseModel):
+    file: Path = Field(description="Absolute path to a file to edit")
+    line: int = Field(description="Line number after which to insert the text (0 for beginning of file)")
+    new_string: str = Field(description="Text to insert")
+
+
+class InsertTool(Tool[InsertToolInput, ToolRunOptions, StringToolOutput]):
+    name = "insert"
+    description = """
+    Inserts the specified text at a specific location in a file.
+    Returns error message on failure.
+    """
+    input_schema = InsertToolInput
+
+    def _create_emitter(self) -> Emitter:
+        return Emitter.root().child(
+            namespace=["tool", "text", self.name],
+            creator=self,
+        )
+
+    async def _run(
+        self, tool_input: InsertToolInput, options: ToolRunOptions | None, context: RunContext
+    ) -> StringToolOutput:
+        try:
+            lines = (await asyncio.to_thread(tool_input.file.read_text)).splitlines(keepends=True)
+            lines.insert(tool_input.line, tool_input.new_string + "\n")
+            await asyncio.to_thread(tool_input.file.write_text, "".join(lines))
+        except Exception as e:
+            return StringToolOutput(result=f"Failed to insert text: {e}")
+        return StringToolOutput()
+
+
+class StrReplaceToolInput(BaseModel):
+    file: Path = Field(description="Absolute path to a file to edit")
+    old_string: str = Field(
+        description="Text to replace (must match exactly, including whitespace and indentation)"
+    )
+    new_string: str = Field(description="New text to insert in place of the old text")
+
+
+class StrReplaceTool(Tool[StrReplaceToolInput, ToolRunOptions, StringToolOutput]):
+    name = "str_replace"
+    description = """
+    Replaces a specific string in the specified file with a new string.
+    Returns error message on failure.
+    """
+    input_schema = StrReplaceToolInput
+
+    def _create_emitter(self) -> Emitter:
+        return Emitter.root().child(
+            namespace=["tool", "text", self.name],
+            creator=self,
+        )
+
+    async def _run(
+        self, tool_input: StrReplaceToolInput, options: ToolRunOptions | None, context: RunContext
+    ) -> StringToolOutput:
+        try:
+            content = await asyncio.to_thread(tool_input.file.read_text)
+            await asyncio.to_thread(
+                tool_input.file.write_text, content.replace(tool_input.old_string, tool_input.new_string)
+            )
+        except Exception as e:
+            return StringToolOutput(result=f"Failed to replace text: {e}")
+        return StringToolOutput()
