diff --git a/auto_gpt/commands.py b/auto_gpt/commands.py index 0133c6653..2f9016ba4 100644 --- a/auto_gpt/commands.py +++ b/auto_gpt/commands.py @@ -79,38 +79,30 @@ class CommandRegistry: commands_list = [f"{idx + 1}. {str(cmd)}" for idx, cmd in enumerate(self.commands.values())] return "\n".join(commands_list) - def scan_directory_for_plugins(self, directory: str) -> None: + def import_commands(self, module_name: str) -> None: """ - Scans the specified directory for Python files containing command plugins. + Imports the specified Python module containing command plugins. - For each file in the directory that ends with ".py", this method imports the associated module and registers any - functions or classes that are decorated with the `AUTO_GPT_COMMAND_IDENTIFIER` attribute as `Command` objects. - The registered `Command` objects are then added to the `commands` dictionary of the `CommandRegistry` object. + This method imports the associated module and registers any functions or + classes that are decorated with the `AUTO_GPT_COMMAND_IDENTIFIER` attribute + as `Command` objects. The registered `Command` objects are then added to the + `commands` dictionary of the `CommandRegistry` object. - Args: - directory (str): The directory to scan for command plugins. - """ + Args: + module_name (str): The name of the module to import for command plugins. + """ - for file in os.listdir(directory): - if file.endswith(".py"): - file_path = os.path.join(directory, file) - module_name = file[:-3] - - spec = importlib.util.spec_from_file_location(module_name, file_path) - module = importlib.util.module_from_spec(spec) - - spec.loader.exec_module(module) - - for attr_name in dir(module): - attr = getattr(module, attr_name) - # Register decorated functions - if hasattr(attr, AUTO_GPT_COMMAND_IDENTIFIER) and getattr(attr, AUTO_GPT_COMMAND_IDENTIFIER): - self.register(attr.command) - # Register command classes - elif inspect.isclass(attr) and issubclass(attr, Command) and attr != Command: - cmd_instance = attr() - self.register(cmd_instance) + module = importlib.import_module(module_name) + for attr_name in dir(module): + attr = getattr(module, attr_name) + # Register decorated functions + if hasattr(attr, AUTO_GPT_COMMAND_IDENTIFIER) and getattr(attr, AUTO_GPT_COMMAND_IDENTIFIER): + self.register(attr.command) + # Register command classes + elif inspect.isclass(attr) and issubclass(attr, Command) and attr != Command: + cmd_instance = attr() + self.register(cmd_instance) def command(name: str, description: str, signature: str = None) -> Callable[..., Any]: """The command decorator is used to create Command objects from ordinary functions.""" diff --git a/auto_gpt/tests/test_commands.py b/auto_gpt/tests/test_commands.py index fc0ccb3d4..a7778b6eb 100644 --- a/auto_gpt/tests/test_commands.py +++ b/auto_gpt/tests/test_commands.py @@ -1,4 +1,5 @@ import shutil +import sys from pathlib import Path import pytest @@ -112,21 +113,19 @@ class TestCommandRegistry: assert f"(arg1: int, arg2: str)" in command_prompt - def test_scan_directory_for_mock_commands(self): - """Test that the registry can scan a directory for mocks command plugins.""" + def test_import_mock_commands_module(self): + """Test that the registry can import a module with mock command plugins.""" registry = CommandRegistry() - mock_commands_dir = Path("/app/auto_gpt/tests/mocks") - import os + mock_commands_module = "auto_gpt.tests.mocks.mock_commands" - print(os.getcwd()) - registry.scan_directory_for_plugins(mock_commands_dir) + registry.import_commands(mock_commands_module) assert "function_based" in registry.commands assert registry.commands["function_based"].name == "function_based" assert registry.commands["function_based"].description == "Function-based test command" - def test_scan_directory_for_temp_command_file(self, tmp_path): - """Test that the registry can scan a directory for command plugins in a temp file.""" + def test_import_temp_command_file_module(self, tmp_path): + """Test that the registry can import a command plugins module from a temp file.""" registry = CommandRegistry() # Create a temp command file @@ -134,8 +133,14 @@ class TestCommandRegistry: temp_commands_file = tmp_path / "mock_commands.py" shutil.copyfile(src, temp_commands_file) - registry.scan_directory_for_plugins(tmp_path) - print(registry.commands) + # Add the temp directory to sys.path to make the module importable + sys.path.append(str(tmp_path)) + + temp_commands_module = "mock_commands" + registry.import_commands(temp_commands_module) + + # Remove the temp directory from sys.path + sys.path.remove(str(tmp_path)) assert "function_based" in registry.commands assert registry.commands["function_based"].name == "function_based" diff --git a/scripts/__init__.py b/scripts/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/scripts/main.py b/scripts/main.py index 5337ceeb8..4d7faa51b 100644 --- a/scripts/main.py +++ b/scripts/main.py @@ -290,7 +290,11 @@ print('Using memory of type: ' + memory.__class__.__name__) # Create a CommandRegistry instance and scan default folder command_registry = CommandRegistry() -command_registry.scan_directory_for_plugins('./scripts') +command_registry.import_commands('scripts.ai_functions') +command_registry.import_commands('scripts.commands') +command_registry.import_commands('scripts.execute_code') +command_registry.import_commands('scripts.agent_manager') +command_registry.import_commands('scripts.file_operations') # Interaction Loop while True: