diff --git a/autogpt/commands/execute_code.py b/autogpt/commands/execute_code.py index 999e40f8..b164a85f 100644 --- a/autogpt/commands/execute_code.py +++ b/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": "", "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": ""') def execute_python_file(filename: str, config: Config) -> str: """Execute a Python file in a Docker container and return the output diff --git a/tests/integration/test_execute_code.py b/tests/integration/test_execute_code.py index ce3c1e16..e4ecf991 100644 --- a/tests/integration/test_execute_code.py +++ b/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()