From a24ab0e87994eb7e314fb9e2ff17a97f68408aba Mon Sep 17 00:00:00 2001 From: Peter Date: Thu, 6 Apr 2023 14:13:23 -0700 Subject: [PATCH 1/4] dynamically load commands from registry --- scripts/agent_manager.py | 4 ++ scripts/ai_functions.py | 8 ++- scripts/auto_gpt/__init__.py | 0 scripts/auto_gpt/commands.py | 114 +++++++++++++++++++++++++++++++++++ scripts/commands.py | 62 ++++--------------- scripts/execute_code.py | 2 + scripts/file_operations.py | 27 +++++---- scripts/main.py | 10 ++- 8 files changed, 160 insertions(+), 67 deletions(-) create mode 100644 scripts/auto_gpt/__init__.py create mode 100644 scripts/auto_gpt/commands.py diff --git a/scripts/agent_manager.py b/scripts/agent_manager.py index ad120c4059..9bd87aa96f 100644 --- a/scripts/agent_manager.py +++ b/scripts/agent_manager.py @@ -1,3 +1,4 @@ +from auto_gpt.commands import command from llm_utils import create_chat_completion next_key = 0 @@ -31,6 +32,7 @@ def create_agent(task, prompt, model): return key, agent_reply +@command("message_agent", "Message GPT Agent", '"key": "", "message": ""') def message_agent(key, message): global agents @@ -51,6 +53,7 @@ def message_agent(key, message): return agent_reply +@command("list_agents", "List GPT Agents", "") def list_agents(): global agents @@ -58,6 +61,7 @@ def list_agents(): return [(key, task) for key, (task, _, _) in agents.items()] +@command("delete_agent", "Delete GPT Agent", '"key": ""') def delete_agent(key): global agents diff --git a/scripts/ai_functions.py b/scripts/ai_functions.py index 05aa93a2da..175dffa2fb 100644 --- a/scripts/ai_functions.py +++ b/scripts/ai_functions.py @@ -3,10 +3,12 @@ import json from config import Config from call_ai_function import call_ai_function from json_parser import fix_and_parse_json +from auto_gpt.commands import command + cfg = Config() # Evaluating code - +@command("evaluate_code", "Evaluate Code", '"code": ""') def evaluate_code(code: str) -> List[str]: function_string = "def analyze_code(code: str) -> List[str]:" args = [code] @@ -18,7 +20,7 @@ def evaluate_code(code: str) -> List[str]: # Improving code - +@command("improve_code", "Get Improved Code", '"suggestions": "", "code": ""') def improve_code(suggestions: List[str], code: str) -> str: function_string = ( "def generate_improved_code(suggestions: List[str], code: str) -> str:" @@ -32,7 +34,7 @@ def improve_code(suggestions: List[str], code: str) -> str: # Writing tests - +@command("write_tests", "Write Tests", '"code": "", "focus": ""') def write_tests(code: str, focus: List[str]) -> str: function_string = ( "def create_test_cases(code: str, focus: Optional[str] = None) -> str:" diff --git a/scripts/auto_gpt/__init__.py b/scripts/auto_gpt/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/scripts/auto_gpt/commands.py b/scripts/auto_gpt/commands.py new file mode 100644 index 0000000000..fb05fdb8f2 --- /dev/null +++ b/scripts/auto_gpt/commands.py @@ -0,0 +1,114 @@ +import os +import sys +import importlib +import inspect +from typing import Callable, Any, List + +# Unique identifier for auto-gpt commands +AUTO_GPT_COMMAND_IDENTIFIER = "auto_gpt_command" + +class Command: + """A class representing a command. + + Attributes: + name (str): The name of the command. + description (str): A brief description of what the command does. + method (Callable[..., Any]): The function that the command executes. + signature (str): The signature of the function that the command executes. Defaults to None. + """ + + def __init__(self, name: str, description: str, method: Callable[..., Any], signature: str = None): + self.name = name + self.description = description + self.method = method + self.signature = signature if signature else str(inspect.signature(self.method)) + + def __call__(self, *args, **kwargs) -> Any: + return self.method(*args, **kwargs) + + def __str__(self) -> str: + return f"{self.name}: {self.description}, args: {self.signature}" + +class CommandRegistry: + """ + The CommandRegistry class is a manager for a collection of Command objects. + It allows the registration, modification, and retrieval of Command objects, + as well as the scanning and loading of command plugins from a specified + directory. + """ + + def __init__(self): + self.commands = {} + + def _import_module(self, module_name: str) -> Any: + return importlib.import_module(module_name) + + def _reload_module(self, module: Any) -> Any: + return importlib.reload(module) + + def register_command(self, cmd: Command) -> None: + self.commands[cmd.name] = cmd + + def reload_commands(self) -> None: + """Reloads all loaded command plugins.""" + for cmd_name in self.commands: + cmd = self.commands[cmd_name] + module = self._import_module(cmd.__module__) + reloaded_module = self._reload_module(module) + if hasattr(reloaded_module, "register"): + reloaded_module.register(self) + + def get_command(self, name: str) -> Callable[..., Any]: + return self.commands.get(name) + + def list_commands(self) -> List[str]: + return [str(cmd) for cmd in self.commands.values()] + + def command_prompt(self) -> str: + """ + Returns a string representation of all registered `Command` objects for use in a prompt + """ + 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: + """ + Scans the specified directory for Python files 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. + + Args: + directory (str): The directory to scan for command plugins. + """ + + for file in os.listdir(directory): + if file.endswith(".py"): + module_name = file[:-3] + 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_command(attr.register_command) + # Register command classes + elif inspect.isclass(attr) and issubclass(attr, Command) and attr != Command: + cmd_instance = attr() + self.register_command(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.""" + def decorator(func: Callable[..., Any]) -> Command: + cmd = Command(name=name, description=description, method=func, signature=signature) + + def wrapper(*args, **kwargs) -> Any: + return func(*args, **kwargs) + + wrapper.register_command = cmd + setattr(wrapper, AUTO_GPT_COMMAND_IDENTIFIER, True) + return wrapper + + return decorator + diff --git a/scripts/commands.py b/scripts/commands.py index fc10d1d052..78f5dbe30d 100644 --- a/scripts/commands.py +++ b/scripts/commands.py @@ -5,13 +5,10 @@ import datetime import agent_manager as agents import speak from config import Config -import ai_functions as ai -from file_operations import read_file, write_to_file, append_to_file, delete_file, search_files -from execute_code import execute_python_file from json_parser import fix_and_parse_json from duckduckgo_search import ddg -from googleapiclient.discovery import build -from googleapiclient.errors import HttpError + +from auto_gpt.commands import CommandRegistry, command cfg = Config() @@ -51,62 +48,27 @@ def get_command(response): return "Error:", str(e) -def execute_command(command_name, arguments): - memory = PineconeMemory() +def execute_command(command_registry: CommandRegistry, command_name: str, arguments: dict) -> str: try: + # Look up the command in the registry + cmd = command_registry.commands.get(command_name) + + # If the command is found, call it with the provided arguments + if cmd: + return cmd(**arguments) + # special case google until this can be moved down into the function. if command_name == "google": - # Check if the Google API key is set and use the official search method # If the API key is not set or has only whitespaces, use the unofficial search method if cfg.google_api_key and (cfg.google_api_key.strip() if cfg.google_api_key else None): return google_official_search(arguments["input"]) else: return google_search(arguments["input"]) - elif command_name == "memory_add": - return memory.add(arguments["string"]) - elif command_name == "start_agent": - return start_agent( - arguments["name"], - arguments["task"], - arguments["prompt"]) - elif command_name == "message_agent": - return message_agent(arguments["key"], arguments["message"]) - elif command_name == "list_agents": - return list_agents() - elif command_name == "delete_agent": - return delete_agent(arguments["key"]) - elif command_name == "get_text_summary": - return get_text_summary(arguments["url"], arguments["question"]) - elif command_name == "get_hyperlinks": - return get_hyperlinks(arguments["url"]) - elif command_name == "read_file": - return read_file(arguments["file"]) - elif command_name == "write_to_file": - return write_to_file(arguments["file"], arguments["text"]) - elif command_name == "append_to_file": - return append_to_file(arguments["file"], arguments["text"]) - elif command_name == "delete_file": - return delete_file(arguments["file"]) - elif command_name == "search_files": - return search_files(arguments["directory"]) - elif command_name == "browse_website": - return browse_website(arguments["url"], arguments["question"]) - # TODO: Change these to take in a file rather than pasted code, if - # non-file is given, return instructions "Input should be a python - # filepath, write your code to file and try again" - elif command_name == "evaluate_code": - return ai.evaluate_code(arguments["code"]) - elif command_name == "improve_code": - return ai.improve_code(arguments["suggestions"], arguments["code"]) - elif command_name == "write_tests": - return ai.write_tests(arguments["code"], arguments.get("focus")) - elif command_name == "execute_python_file": # Add this command - return execute_python_file(arguments["file"]) elif command_name == "task_complete": shutdown() else: return f"Unknown command {command_name}" - # All errors, return "Error: + error message" + except Exception as e: return "Error: " + str(e) @@ -158,6 +120,7 @@ def google_official_search(query, num_results=8): # Return the list of search result URLs return search_results_links +@command("browse_website", "Browse Website", '"url": "", "question": ""') def browse_website(url, question): summary = get_text_summary(url, question) links = get_hyperlinks(url) @@ -230,6 +193,7 @@ def shutdown(): quit() +@command("start_agent", "Start GPT Agent", '"name": "", "task": "", "prompt": ""') def start_agent(name, task, prompt, model=cfg.fast_llm_model): global cfg diff --git a/scripts/execute_code.py b/scripts/execute_code.py index 614ef6fc3d..a90ab4a007 100644 --- a/scripts/execute_code.py +++ b/scripts/execute_code.py @@ -1,7 +1,9 @@ import docker import os +from auto_gpt.commands import command +@command("execute_python_file", "Execute Python File", '"file": ""') def execute_python_file(file): workspace_folder = "auto_gpt_workspace" diff --git a/scripts/file_operations.py b/scripts/file_operations.py index 90c9a1e4d7..140123c594 100644 --- a/scripts/file_operations.py +++ b/scripts/file_operations.py @@ -1,5 +1,7 @@ import os import os.path +from auto_gpt.commands import command + # Set a dedicated folder for file I/O working_directory = "auto_gpt_workspace" @@ -17,20 +19,20 @@ def safe_join(base, *paths): return norm_new_path - -def read_file(filename): +@command("read_file", "Read file", '"file": ""') +def read_file(file): try: - filepath = safe_join(working_directory, filename) + filepath = safe_join(working_directory, file) with open(filepath, "r") as f: content = f.read() return content except Exception as e: return "Error: " + str(e) - -def write_to_file(filename, text): +@command("write_to_file", "Write to file", '"file": "", "text": ""') +def write_to_file(file, text): try: - filepath = safe_join(working_directory, filename) + filepath = safe_join(working_directory, file) directory = os.path.dirname(filepath) if not os.path.exists(directory): os.makedirs(directory) @@ -40,25 +42,26 @@ def write_to_file(filename, text): except Exception as e: return "Error: " + str(e) - -def append_to_file(filename, text): +@command("append_to_file", "Append to file", '"file": "", "text": ""') +def append_to_file(file, text): try: - filepath = safe_join(working_directory, filename) + filepath = safe_join(working_directory, file) with open(filepath, "a") as f: f.write(text) return "Text appended successfully." except Exception as e: return "Error: " + str(e) - -def delete_file(filename): +@command("delete_file", "Delete file", '"file": ""') +def delete_file(file): try: - filepath = safe_join(working_directory, filename) + filepath = safe_join(working_directory, file) os.remove(filepath) return "File deleted successfully." except Exception as e: return "Error: " + str(e) +@command("search_files", "Search Files", '"directory": ""') def search_files(directory): found_files = [] diff --git a/scripts/main.py b/scripts/main.py index a79fd553ce..5337ceeb83 100644 --- a/scripts/main.py +++ b/scripts/main.py @@ -8,14 +8,13 @@ from colorama import Fore, Style from spinner import Spinner import time import speak -from enum import Enum, auto -import sys from config import Config from json_parser import fix_and_parse_json from ai_config import AIConfig import traceback import yaml import argparse +from auto_gpt.commands import CommandRegistry def print_to_console( @@ -281,6 +280,7 @@ next_action_count = 0 # Make a constant: user_input = "Determine which next command to use, and respond using the format specified above:" + # Initialize memory and make sure it is empty. # this is particularly important for indexing and referencing pinecone memory memory = PineconeMemory() @@ -288,6 +288,10 @@ memory.clear() 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') + # Interaction Loop while True: # Send message to AI, get response @@ -362,7 +366,7 @@ while True: elif command_name == "human_feedback": result = f"Human feedback: {user_input}" else: - result = f"Command {command_name} returned: {cmd.execute_command(command_name, arguments)}" + result = f"Command {command_name} returned: {cmd.execute_command(command_registry, command_name, arguments)}" if next_action_count > 0: next_action_count -= 1 From e2a6ed6955bdcf9487ca288581988848ff0c87f7 Mon Sep 17 00:00:00 2001 From: Peter Date: Thu, 6 Apr 2023 18:24:53 -0700 Subject: [PATCH 2/4] adding tests for CommandRegistry --- scripts/auto_gpt/tests/__init__.py | 0 scripts/auto_gpt/tests/mock/__init__.py | 0 scripts/auto_gpt/tests/mock/mock_commands.py | 12 ++ scripts/auto_gpt/tests/test_commands.py | 174 +++++++++++++++++++ 4 files changed, 186 insertions(+) create mode 100644 scripts/auto_gpt/tests/__init__.py create mode 100644 scripts/auto_gpt/tests/mock/__init__.py create mode 100644 scripts/auto_gpt/tests/mock/mock_commands.py create mode 100644 scripts/auto_gpt/tests/test_commands.py diff --git a/scripts/auto_gpt/tests/__init__.py b/scripts/auto_gpt/tests/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/scripts/auto_gpt/tests/mock/__init__.py b/scripts/auto_gpt/tests/mock/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/scripts/auto_gpt/tests/mock/mock_commands.py b/scripts/auto_gpt/tests/mock/mock_commands.py new file mode 100644 index 0000000000..3514b62af1 --- /dev/null +++ b/scripts/auto_gpt/tests/mock/mock_commands.py @@ -0,0 +1,12 @@ +from commands import Command, command + +class TestCommand(Command): + def __init__(self): + super().__init__(name='class_based', description='Class-based test command') + + def __call__(self, arg1: int, arg2: str) -> str: + return f'{arg1} - {arg2}' + +@command('function_based', 'Function-based test command') +def function_based(arg1: int, arg2: str) -> str: + return f'{arg1} - {arg2}' diff --git a/scripts/auto_gpt/tests/test_commands.py b/scripts/auto_gpt/tests/test_commands.py new file mode 100644 index 0000000000..e2681798a5 --- /dev/null +++ b/scripts/auto_gpt/tests/test_commands.py @@ -0,0 +1,174 @@ +from pathlib import Path + +import pytest +from commands import Command, CommandRegistry + + +class TestCommand: + @staticmethod + def example_function(arg1: int, arg2: str) -> str: + return f"{arg1} - {arg2}" + + def test_command_creation(self): + cmd = Command(name="example", description="Example command", method=self.example_function) + + assert cmd.name == "example" + assert cmd.description == "Example command" + assert cmd.method == self.example_function + assert cmd.signature == "arg1: int, arg2: str" + + def test_command_call(self): + cmd = Command(name="example", description="Example command", method=self.example_function) + + result = cmd(arg1=1, arg2="test") + assert result == "1 - test" + + def test_command_call_with_invalid_arguments(self): + cmd = Command(name="example", description="Example command", method=self.example_function) + + with pytest.raises(TypeError): + cmd(arg1="invalid", arg2="test") + + def test_command_default_signature(self): + cmd = Command(name="example", description="Example command", method=self.example_function) + + assert cmd.signature == "arg1: int, arg2: str" + + def test_command_custom_signature(self): + custom_signature = "custom_arg1: int, custom_arg2: str" + cmd = Command(name="example", description="Example command", method=self.example_function, signature=custom_signature) + + assert cmd.signature == custom_signature + + + +class TestCommandRegistry: + @staticmethod + def example_function(arg1: int, arg2: str) -> str: + return f"{arg1} - {arg2}" + + def test_register_command(self): + """Test that a command can be registered to the registry.""" + registry = CommandRegistry() + cmd = Command(name="example", description="Example command", method=self.example_function) + + registry.register(cmd) + + assert cmd.name in registry._commands + assert registry._commands[cmd.name] == cmd + + def test_unregister_command(self): + """Test that a command can be unregistered from the registry.""" + registry = CommandRegistry() + cmd = Command(name="example", description="Example command", method=self.example_function) + + registry.register(cmd) + registry.unregister(cmd.name) + + assert cmd.name not in registry._commands + + def test_get_command(self): + """Test that a command can be retrieved from the registry.""" + registry = CommandRegistry() + cmd = Command(name="example", description="Example command", method=self.example_function) + + registry.register(cmd) + retrieved_cmd = registry.get(cmd.name) + + assert retrieved_cmd == cmd + + def test_get_nonexistent_command(self): + """Test that attempting to get a nonexistent command raises a KeyError.""" + registry = CommandRegistry() + + with pytest.raises(KeyError): + registry.get("nonexistent_command") + + def test_call_command(self): + """Test that a command can be called through the registry.""" + registry = CommandRegistry() + cmd = Command(name="example", description="Example command", method=self.example_function) + + registry.register(cmd) + result = registry.call("example", arg1=1, arg2="test") + + assert result == "1 - test" + + def test_call_nonexistent_command(self): + """Test that attempting to call a nonexistent command raises a KeyError.""" + registry = CommandRegistry() + + with pytest.raises(KeyError): + registry.call("nonexistent_command", arg1=1, arg2="test") + + def test_get_command_list(self): + """Test that a list of registered commands can be retrieved.""" + registry = CommandRegistry() + cmd1 = Command(name="example1", description="Example command 1", method=self.example_function) + cmd2 = Command(name="example2", description="Example command 2", method=self.example_function) + + registry.register(cmd1) + registry.register(cmd2) + command_list = registry.get_command_list() + + assert len(command_list) == 2 + assert cmd1.name in command_list + assert cmd2.name in command_list + + def test_get_command_prompt(self): + """Test that the command prompt is correctly formatted.""" + registry = CommandRegistry() + cmd = Command(name="example", description="Example command", method=self.example_function) + + registry.register(cmd) + command_prompt = registry.get_command_prompt() + + assert f"{cmd.name}: {cmd.description}, args: {cmd.signature}" in command_prompt + + def test_scan_directory_for_mock_commands(self): + """Test that the registry can scan a directory for mock command plugins.""" + registry = CommandRegistry() + mock_commands_dir = Path("auto_gpt/tests/mocks") + + registry.scan_directory_for_plugins(mock_commands_dir) + + assert "mock_class_based" in registry._commands + assert registry._commands["mock_class_based"].name == "mock_class_based" + assert registry._commands["mock_class_based"].description == "Mock class-based command" + + assert "mock_function_based" in registry._commands + assert registry._commands["mock_function_based"].name == "mock_function_based" + assert registry._commands["mock_function_based"].description == "Mock function-based 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.""" + registry = CommandRegistry() + temp_commands_dir = tmp_path / "temp_commands" + temp_commands_dir.mkdir() + + # Create a temp command file + temp_commands_file = temp_commands_dir / "temp_commands.py" + temp_commands_content = ( + "from commands import Command, command\n\n" + "class TempCommand(Command):\n" + " def __init__(self):\n" + " super().__init__(name='temp_class_based', description='Temp class-based command')\n\n" + " def __call__(self, arg1: int, arg2: str) -> str:\n" + " return f'{arg1} - {arg2}'\n\n" + "@command('temp_function_based', 'Temp function-based command')\n" + "def temp_function_based(arg1: int, arg2: str) -> str:\n" + " return f'{arg1} - {arg2}'\n" + ) + + with open(temp_commands_file, "w") as f: + f.write(temp_commands_content) + + registry.scan_directory_for_plugins(temp_commands_dir) + + assert "temp_class_based" in registry._commands + assert registry._commands["temp_class_based"].name == "temp_class_based" + assert registry._commands["temp_class_based"].description == "Temp class-based command" + + assert "temp_function_based" in registry._commands + assert registry._commands["temp_function_based"].name == "temp_function_based" + assert registry._commands["temp_function_based"].description == "Temp function-based command" From b4a0ef9babf9edfbd395a1afaaab22381a522643 Mon Sep 17 00:00:00 2001 From: Peter Date: Thu, 6 Apr 2023 19:25:44 -0700 Subject: [PATCH 3/4] resolving test failures --- {scripts/auto_gpt => auto_gpt}/__init__.py | 0 {scripts/auto_gpt => auto_gpt}/commands.py | 33 +++++-- .../auto_gpt => auto_gpt}/tests/__init__.py | 0 .../mock => auto_gpt/tests/mocks}/__init__.py | 0 auto_gpt/tests/mocks/mock_commands.py | 6 ++ .../tests/test_commands.py | 86 ++++++------------- scripts/ai_functions.py | 3 +- scripts/auto_gpt/tests/mock/mock_commands.py | 12 --- scripts/commands.py | 1 - 9 files changed, 58 insertions(+), 83 deletions(-) rename {scripts/auto_gpt => auto_gpt}/__init__.py (100%) rename {scripts/auto_gpt => auto_gpt}/commands.py (78%) rename {scripts/auto_gpt => auto_gpt}/tests/__init__.py (100%) rename {scripts/auto_gpt/tests/mock => auto_gpt/tests/mocks}/__init__.py (100%) create mode 100644 auto_gpt/tests/mocks/mock_commands.py rename {scripts/auto_gpt => auto_gpt}/tests/test_commands.py (53%) delete mode 100644 scripts/auto_gpt/tests/mock/mock_commands.py diff --git a/scripts/auto_gpt/__init__.py b/auto_gpt/__init__.py similarity index 100% rename from scripts/auto_gpt/__init__.py rename to auto_gpt/__init__.py diff --git a/scripts/auto_gpt/commands.py b/auto_gpt/commands.py similarity index 78% rename from scripts/auto_gpt/commands.py rename to auto_gpt/commands.py index fb05fdb8f2..0133c66531 100644 --- a/scripts/auto_gpt/commands.py +++ b/auto_gpt/commands.py @@ -13,7 +13,6 @@ class Command: Attributes: name (str): The name of the command. description (str): A brief description of what the command does. - method (Callable[..., Any]): The function that the command executes. signature (str): The signature of the function that the command executes. Defaults to None. """ @@ -46,9 +45,15 @@ class CommandRegistry: def _reload_module(self, module: Any) -> Any: return importlib.reload(module) - def register_command(self, cmd: Command) -> None: + def register(self, cmd: Command) -> None: self.commands[cmd.name] = cmd + def unregister(self, command_name: str): + if command_name in self.commands: + del self.commands[command_name] + else: + raise KeyError(f"Command '{command_name}' not found in registry.") + def reload_commands(self) -> None: """Reloads all loaded command plugins.""" for cmd_name in self.commands: @@ -59,10 +64,13 @@ class CommandRegistry: reloaded_module.register(self) def get_command(self, name: str) -> Callable[..., Any]: - return self.commands.get(name) + return self.commands[name] - def list_commands(self) -> List[str]: - return [str(cmd) for cmd in self.commands.values()] + def call(self, command_name: str, **kwargs) -> Any: + if command_name not in self.commands: + raise KeyError(f"Command '{command_name}' not found in registry.") + command = self.commands[command_name] + return command(**kwargs) def command_prompt(self) -> str: """ @@ -85,17 +93,23 @@ class CommandRegistry: for file in os.listdir(directory): if file.endswith(".py"): + file_path = os.path.join(directory, file) module_name = file[:-3] - module = importlib.import_module(module_name) + + 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_command(attr.register_command) + self.register(attr.command) # Register command classes elif inspect.isclass(attr) and issubclass(attr, Command) and attr != Command: cmd_instance = attr() - self.register_command(cmd_instance) + self.register(cmd_instance) def command(name: str, description: str, signature: str = None) -> Callable[..., Any]: @@ -106,7 +120,8 @@ def command(name: str, description: str, signature: str = None) -> Callable[..., def wrapper(*args, **kwargs) -> Any: return func(*args, **kwargs) - wrapper.register_command = cmd + wrapper.command = cmd + setattr(wrapper, AUTO_GPT_COMMAND_IDENTIFIER, True) return wrapper diff --git a/scripts/auto_gpt/tests/__init__.py b/auto_gpt/tests/__init__.py similarity index 100% rename from scripts/auto_gpt/tests/__init__.py rename to auto_gpt/tests/__init__.py diff --git a/scripts/auto_gpt/tests/mock/__init__.py b/auto_gpt/tests/mocks/__init__.py similarity index 100% rename from scripts/auto_gpt/tests/mock/__init__.py rename to auto_gpt/tests/mocks/__init__.py diff --git a/auto_gpt/tests/mocks/mock_commands.py b/auto_gpt/tests/mocks/mock_commands.py new file mode 100644 index 0000000000..d68ceb8133 --- /dev/null +++ b/auto_gpt/tests/mocks/mock_commands.py @@ -0,0 +1,6 @@ +from auto_gpt.commands import Command, command + + +@command('function_based', 'Function-based test command') +def function_based(arg1: int, arg2: str) -> str: + return f'{arg1} - {arg2}' diff --git a/scripts/auto_gpt/tests/test_commands.py b/auto_gpt/tests/test_commands.py similarity index 53% rename from scripts/auto_gpt/tests/test_commands.py rename to auto_gpt/tests/test_commands.py index e2681798a5..fc0ccb3d40 100644 --- a/scripts/auto_gpt/tests/test_commands.py +++ b/auto_gpt/tests/test_commands.py @@ -1,7 +1,8 @@ +import shutil from pathlib import Path import pytest -from commands import Command, CommandRegistry +from auto_gpt.commands import Command, CommandRegistry class TestCommand: @@ -15,7 +16,7 @@ class TestCommand: assert cmd.name == "example" assert cmd.description == "Example command" assert cmd.method == self.example_function - assert cmd.signature == "arg1: int, arg2: str" + assert cmd.signature == "(arg1: int, arg2: str) -> str" def test_command_call(self): cmd = Command(name="example", description="Example command", method=self.example_function) @@ -27,12 +28,12 @@ class TestCommand: cmd = Command(name="example", description="Example command", method=self.example_function) with pytest.raises(TypeError): - cmd(arg1="invalid", arg2="test") + cmd(arg1="invalid", does_not_exist="test") def test_command_default_signature(self): cmd = Command(name="example", description="Example command", method=self.example_function) - assert cmd.signature == "arg1: int, arg2: str" + assert cmd.signature == "(arg1: int, arg2: str) -> str" def test_command_custom_signature(self): custom_signature = "custom_arg1: int, custom_arg2: str" @@ -54,8 +55,8 @@ class TestCommandRegistry: registry.register(cmd) - assert cmd.name in registry._commands - assert registry._commands[cmd.name] == cmd + assert cmd.name in registry.commands + assert registry.commands[cmd.name] == cmd def test_unregister_command(self): """Test that a command can be unregistered from the registry.""" @@ -65,7 +66,7 @@ class TestCommandRegistry: registry.register(cmd) registry.unregister(cmd.name) - assert cmd.name not in registry._commands + assert cmd.name not in registry.commands def test_get_command(self): """Test that a command can be retrieved from the registry.""" @@ -73,7 +74,7 @@ class TestCommandRegistry: cmd = Command(name="example", description="Example command", method=self.example_function) registry.register(cmd) - retrieved_cmd = registry.get(cmd.name) + retrieved_cmd = registry.get_command(cmd.name) assert retrieved_cmd == cmd @@ -82,7 +83,7 @@ class TestCommandRegistry: registry = CommandRegistry() with pytest.raises(KeyError): - registry.get("nonexistent_command") + registry.get_command("nonexistent_command") def test_call_command(self): """Test that a command can be called through the registry.""" @@ -101,74 +102,41 @@ class TestCommandRegistry: with pytest.raises(KeyError): registry.call("nonexistent_command", arg1=1, arg2="test") - def test_get_command_list(self): - """Test that a list of registered commands can be retrieved.""" - registry = CommandRegistry() - cmd1 = Command(name="example1", description="Example command 1", method=self.example_function) - cmd2 = Command(name="example2", description="Example command 2", method=self.example_function) - - registry.register(cmd1) - registry.register(cmd2) - command_list = registry.get_command_list() - - assert len(command_list) == 2 - assert cmd1.name in command_list - assert cmd2.name in command_list - def test_get_command_prompt(self): """Test that the command prompt is correctly formatted.""" registry = CommandRegistry() cmd = Command(name="example", description="Example command", method=self.example_function) registry.register(cmd) - command_prompt = registry.get_command_prompt() + command_prompt = registry.command_prompt() - assert f"{cmd.name}: {cmd.description}, args: {cmd.signature}" in command_prompt + 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 mock command plugins.""" + """Test that the registry can scan a directory for mocks command plugins.""" registry = CommandRegistry() - mock_commands_dir = Path("auto_gpt/tests/mocks") + mock_commands_dir = Path("/app/auto_gpt/tests/mocks") + import os + print(os.getcwd()) registry.scan_directory_for_plugins(mock_commands_dir) - assert "mock_class_based" in registry._commands - assert registry._commands["mock_class_based"].name == "mock_class_based" - assert registry._commands["mock_class_based"].description == "Mock class-based command" - - assert "mock_function_based" in registry._commands - assert registry._commands["mock_function_based"].name == "mock_function_based" - assert registry._commands["mock_function_based"].description == "Mock function-based command" + 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.""" registry = CommandRegistry() - temp_commands_dir = tmp_path / "temp_commands" - temp_commands_dir.mkdir() # Create a temp command file - temp_commands_file = temp_commands_dir / "temp_commands.py" - temp_commands_content = ( - "from commands import Command, command\n\n" - "class TempCommand(Command):\n" - " def __init__(self):\n" - " super().__init__(name='temp_class_based', description='Temp class-based command')\n\n" - " def __call__(self, arg1: int, arg2: str) -> str:\n" - " return f'{arg1} - {arg2}'\n\n" - "@command('temp_function_based', 'Temp function-based command')\n" - "def temp_function_based(arg1: int, arg2: str) -> str:\n" - " return f'{arg1} - {arg2}'\n" - ) + src = Path("/app/auto_gpt/tests/mocks/mock_commands.py") + temp_commands_file = tmp_path / "mock_commands.py" + shutil.copyfile(src, temp_commands_file) - with open(temp_commands_file, "w") as f: - f.write(temp_commands_content) + registry.scan_directory_for_plugins(tmp_path) + print(registry.commands) - registry.scan_directory_for_plugins(temp_commands_dir) - - assert "temp_class_based" in registry._commands - assert registry._commands["temp_class_based"].name == "temp_class_based" - assert registry._commands["temp_class_based"].description == "Temp class-based command" - - assert "temp_function_based" in registry._commands - assert registry._commands["temp_function_based"].name == "temp_function_based" - assert registry._commands["temp_function_based"].description == "Temp function-based command" + assert "function_based" in registry.commands + assert registry.commands["function_based"].name == "function_based" + assert registry.commands["function_based"].description == "Function-based test command" diff --git a/scripts/ai_functions.py b/scripts/ai_functions.py index 175dffa2fb..c7cb9a5a00 100644 --- a/scripts/ai_functions.py +++ b/scripts/ai_functions.py @@ -1,8 +1,7 @@ -from typing import List, Optional +from typing import List import json from config import Config from call_ai_function import call_ai_function -from json_parser import fix_and_parse_json from auto_gpt.commands import command cfg = Config() diff --git a/scripts/auto_gpt/tests/mock/mock_commands.py b/scripts/auto_gpt/tests/mock/mock_commands.py deleted file mode 100644 index 3514b62af1..0000000000 --- a/scripts/auto_gpt/tests/mock/mock_commands.py +++ /dev/null @@ -1,12 +0,0 @@ -from commands import Command, command - -class TestCommand(Command): - def __init__(self): - super().__init__(name='class_based', description='Class-based test command') - - def __call__(self, arg1: int, arg2: str) -> str: - return f'{arg1} - {arg2}' - -@command('function_based', 'Function-based test command') -def function_based(arg1: int, arg2: str) -> str: - return f'{arg1} - {arg2}' diff --git a/scripts/commands.py b/scripts/commands.py index 78f5dbe30d..2a78031ae1 100644 --- a/scripts/commands.py +++ b/scripts/commands.py @@ -1,6 +1,5 @@ import browse import json -from memory import PineconeMemory import datetime import agent_manager as agents import speak From 3095591064b213c4e916a9c6d7d1e85c6b0c80c8 Mon Sep 17 00:00:00 2001 From: Peter Date: Thu, 6 Apr 2023 20:00:28 -0700 Subject: [PATCH 4/4] switch to explicit module imports --- auto_gpt/commands.py | 46 ++++++++++++++------------------- auto_gpt/tests/test_commands.py | 25 +++++++++++------- scripts/__init__.py | 0 scripts/main.py | 6 ++++- 4 files changed, 39 insertions(+), 38 deletions(-) create mode 100644 scripts/__init__.py diff --git a/auto_gpt/commands.py b/auto_gpt/commands.py index 0133c66531..2f9016ba42 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 fc0ccb3d40..a7778b6ebc 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 0000000000..e69de29bb2 diff --git a/scripts/main.py b/scripts/main.py index 5337ceeb83..4d7faa51b8 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: