refactor tools/ to have a single run_command

Signed-off-by: Tomas Tomecek <[email protected]>

Author: TomasTomecek

beeai/agents/tools/commands.py

@@ -7,6 +7,8 @@
 from beeai_framework.emitter import Emitter
 from beeai_framework.tools import JSONToolOutput, Tool, ToolRunOptions
 
+from tools.utils import run_command
+
 
 class RunShellCommandToolInput(BaseModel):
     command: str = Field(description="Command to run")
@@ -39,18 +41,6 @@ def _create_emitter(self) -> Emitter:
     async def _run(
         self, tool_input: RunShellCommandToolInput, options: ToolRunOptions | None, context: RunContext
     ) -> RunShellCommandToolOutput:
-        proc = await asyncio.create_subprocess_shell(
-            tool_input.command,
-            stdout=asyncio.subprocess.PIPE,
-            stderr=asyncio.subprocess.PIPE,
-        )
-
-        stdout, stderr = await proc.communicate()
-
-        result = {
-            "exit_code": proc.returncode,
-            "stdout": stdout.decode() if stdout else None,
-            "stderr": stderr.decode() if stderr else None,
-        }
+        result = await run_command(tool_input.command, subprocess_function=asyncio.create_subprocess_shell)
 
         return RunShellCommandToolOutput(RunShellCommandToolResult.model_validate(result))

beeai/agents/tools/utils.py

@@ -0,0 +1,47 @@
+import asyncio
+from pathlib import Path
+from typing import Callable, Optional, Union
+
+
+async def run_command(
+    cmd: Union[list[str], str],
+    cwd: Optional[Path] = None,
+    subprocess_function: Callable = asyncio.create_subprocess_exec
+) -> dict[str, str | int]:
+    """
+    Run a shell command asynchronously and capture its output.
+
+    Args:
+        cmd: The command to run. Can be a list of arguments (for exec) or a string (for shell).
+        cwd: Optional working directory to run the command in.
+        subprocess_function: The asyncio subprocess function to use (e.g., create_subprocess_exec or create_subprocess_shell).
+
+    Returns:
+        A dictionary with:
+            - "exit_code": The process exit code (int)
+            - "stdout": The standard output as a string, or None
+            - "stderr": The standard error as a string, or None
+    """
+    if subprocess_function is asyncio.create_subprocess_exec:
+        proc = await subprocess_function(
+            cmd[0],
+            *cmd[1:],
+            stdout=asyncio.subprocess.PIPE,
+            stderr=asyncio.subprocess.PIPE,
+            cwd=cwd,
+        )
+    else:
+        proc = await subprocess_function(
+            cmd,
+            stdout=asyncio.subprocess.PIPE,
+            stderr=asyncio.subprocess.PIPE,
+            cwd=cwd,
+        )
+
+    stdout, stderr = await proc.communicate()
+
+    return {
+        "exit_code": proc.returncode,
+        "stdout": stdout.decode() if stdout else None,
+        "stderr": stderr.decode() if stderr else None,
+    }

beeai/agents/tools/wicked_git.py

@@ -1,8 +1,9 @@
-import asyncio
 from pathlib import Path
 
 from pydantic import BaseModel, Field
 
+from tools.utils import run_command
+
 from beeai_framework.context import RunContext
 from beeai_framework.emitter import Emitter
 from beeai_framework.tools import StringToolOutput, Tool, ToolRunOptions
@@ -13,23 +14,6 @@ class GitPatchCreationToolInput(BaseModel):
     patch_file_path: str = Field(description="Absolute path where the patch file should be saved")
 
 
-async def run_command(cmd: list[str], cwd: Path) -> dict[str, str | int]:
-    proc = await asyncio.create_subprocess_exec(
-        cmd[0],
-        *cmd[1:],
-        stdout=asyncio.subprocess.PIPE,
-        stderr=asyncio.subprocess.PIPE,
-        cwd=cwd,
-    )
-
-    stdout, stderr = await proc.communicate()
-
-    return {
-        "exit_code": proc.returncode,
-        "stdout": stdout.decode() if stdout else None,
-        "stderr": stderr.decode() if stderr else None,
-    }
-
 class GitPatchCreationTool(Tool[GitPatchCreationToolInput, ToolRunOptions, StringToolOutput]):
     name = "git_patch_create"
     description = """