Add command for directly executing python code (#4581)

* Add command for directly executing python code

* Fix docstring

* Clarify / update filename references

---------

Co-authored-by: merwanehamadi <merwanehamadi@gmail.com>

autogpt/commands/execute_code.py

@@ -8,11 +8,46 @@ from docker.errors import ImageNotFound
 
 from autogpt.commands.command import command
 from autogpt.config import Config
+from autogpt.config.ai_config import AIConfig
 from autogpt.logs import logger
 from autogpt.setup import CFG
 from autogpt.workspace.workspace import Workspace
 
 
+@command(
+    "execute_python_code",
+    "Create a Python file and execute it",
+    '"code": "<code>", "basename": "<basename>"',
+)
+def execute_python_code(code: str, basename: str, config: Config) -> str:
+    """Create and execute a Python file in a Docker container and return the STDOUT of the
+    executed code. If there is any data that needs to be captured use a print statement
+
+    Args:
+        code (str): The Python code to run
+        basename (str): A name to be given to the Python file
+
+    Returns:
+        str: The STDOUT captured from the code when it ran
+    """
+    ai_name = AIConfig.load(config.ai_settings_file).ai_name
+    directory = os.path.join(config.workspace_path, ai_name, "executed_code")
+    os.makedirs(directory, exist_ok=True)
+
+    if not basename.endswith(".py"):
+        basename = basename + ".py"
+
+    path = os.path.join(directory, basename)
+
+    try:
+        with open(path, "w+", encoding="utf-8") as f:
+            f.write(code)
+
+        return execute_python_file(f.name, config)
+    except Exception as e:
+        return f"Error: {str(e)}"
+
+
 @command("execute_python_file", "Execute Python File", '"filename": "<filename>"')
 def execute_python_file(filename: str, config: Config) -> str:
     """Execute a Python file in a Docker container and return the output

tests/integration/test_execute_code.py

@@ -1,3 +1,4 @@
+import os
 import random
 import string
 import tempfile
@@ -8,6 +9,7 @@ from pytest_mock import MockerFixture
 
 import autogpt.commands.execute_code as sut  # system under testing
 from autogpt.config import Config
+from autogpt.config.ai_config import AIConfig
 
 
 @pytest.fixture
@@ -16,9 +18,14 @@ def config_allow_execute(config: Config, mocker: MockerFixture) -> Callable:
 
 
 @pytest.fixture
-def python_test_file(config: Config, random_string) -> Callable:
+def random_code(random_string) -> Callable:
+    return f"print('Hello {random_string}!')"
+
+
+@pytest.fixture
+def python_test_file(config: Config, random_code: str) -> Callable:
     temp_file = tempfile.NamedTemporaryFile(dir=config.workspace_path, suffix=".py")
-    temp_file.write(str.encode(f"print('Hello {random_string}!')"))
+    temp_file.write(str.encode(random_code))
     temp_file.flush()
 
     yield temp_file.name
@@ -35,6 +42,39 @@ def test_execute_python_file(python_test_file: str, random_string: str, config):
     assert result.replace("\r", "") == f"Hello {random_string}!\n"
 
 
+def test_execute_python_code(random_code: str, random_string: str, config: Config):
+    ai_name = AIConfig.load(config.ai_settings_file).ai_name
+
+    result: str = sut.execute_python_code(random_code, "test_code", config)
+    assert result.replace("\r", "") == f"Hello {random_string}!\n"
+
+    # Check that the code is stored
+    destination = os.path.join(
+        config.workspace_path, ai_name, "executed_code", "test_code.py"
+    )
+    with open(destination) as f:
+        assert f.read() == random_code
+
+
+def test_execute_python_code_overwrites_file(
+    random_code: str, random_string: str, config: Config
+):
+    ai_name = AIConfig.load(config.ai_settings_file).ai_name
+    destination = os.path.join(
+        config.workspace_path, ai_name, "executed_code", "test_code.py"
+    )
+    os.makedirs(os.path.dirname(destination), exist_ok=True)
+
+    with open(destination, "w+") as f:
+        f.write("This will be overwritten")
+
+    sut.execute_python_code(random_code, "test_code.py", config)
+
+    # Check that the file is updated with the new code
+    with open(destination) as f:
+        assert f.read() == random_code
+
+
 def test_execute_python_file_invalid(config: Config):
     assert all(
         s in sut.execute_python_file("not_python", config).lower()