diff --git a/auto_gpt/__init__.py b/auto_gpt/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/auto_gpt/commands.py b/auto_gpt/commands.py new file mode 100644 index 0000000000..2f9016ba42 --- /dev/null +++ b/auto_gpt/commands.py @@ -0,0 +1,121 @@ +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. + 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(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: + 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[name] + + 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: + """ + 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 import_commands(self, module_name: str) -> None: + """ + Imports the specified Python module containing command plugins. + + 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: + module_name (str): The name of the module to import for command plugins. + """ + + 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.""" + 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.command = cmd + + setattr(wrapper, AUTO_GPT_COMMAND_IDENTIFIER, True) + return wrapper + + return decorator + diff --git a/auto_gpt/tests/__init__.py b/auto_gpt/tests/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/auto_gpt/tests/mocks/__init__.py b/auto_gpt/tests/mocks/__init__.py new file mode 100644 index 0000000000..e69de29bb2 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/auto_gpt/tests/test_commands.py b/auto_gpt/tests/test_commands.py new file mode 100644 index 0000000000..a7778b6ebc --- /dev/null +++ b/auto_gpt/tests/test_commands.py @@ -0,0 +1,147 @@ +import shutil +import sys +from pathlib import Path + +import pytest +from auto_gpt.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) -> 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", 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) -> 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_command(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_command("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_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.command_prompt() + + assert f"(arg1: int, arg2: str)" in command_prompt + + def test_import_mock_commands_module(self): + """Test that the registry can import a module with mock command plugins.""" + registry = CommandRegistry() + mock_commands_module = "auto_gpt.tests.mocks.mock_commands" + + 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_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 + src = Path("/app/auto_gpt/tests/mocks/mock_commands.py") + temp_commands_file = tmp_path / "mock_commands.py" + shutil.copyfile(src, temp_commands_file) + + # 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" + assert registry.commands["function_based"].description == "Function-based test command" diff --git a/autogpt/__main__.py b/autogpt/__main__.py index 3e3331548a..60079f7511 100644 --- a/autogpt/__main__.py +++ b/autogpt/__main__.py @@ -5,7 +5,7 @@ from pathlib import Path from colorama import Fore from autogpt.agent.agent import Agent from autogpt.args import parse_arguments - +from autogpt.commands.command import CommandRegistry from autogpt.config import Config, check_openai_api_key from autogpt.logs import logger from autogpt.memory import get_memory @@ -45,7 +45,13 @@ def main() -> None: print(f"{plugin._name}: {plugin._version} - {plugin._description}") cfg.set_plugins(loaded_plugins) - + # Create a CommandRegistry instance and scan default folder + command_registry = CommandRegistry() + 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') ai_name = "" ai_config = construct_main_ai_config() # print(prompt) @@ -69,6 +75,7 @@ def main() -> None: memory=memory, full_message_history=full_message_history, next_action_count=next_action_count, + command_registry=command_registry, config=ai_config, prompt=ai_config.construct_full_prompt(), user_input=user_input, diff --git a/autogpt/agent/agent.py b/autogpt/agent/agent.py index a3a3729a4e..513478b9d8 100644 --- a/autogpt/agent/agent.py +++ b/autogpt/agent/agent.py @@ -31,6 +31,7 @@ class Agent: memory, full_message_history, next_action_count, + command_registry, config, prompt, user_input, @@ -39,6 +40,7 @@ class Agent: self.memory = memory self.full_message_history = full_message_history self.next_action_count = next_action_count + self.command_registry = command_registry self.config = config self.prompt = prompt self.user_input = user_input @@ -167,7 +169,7 @@ class Agent: ) result = ( f"Command {command_name} returned: " - f"{execute_command(command_name, arguments, self.config.prompt_generator)}" + f"{execute_command(self.command_registry, command_name, arguments, self.config.prompt_generator)}" ) for plugin in cfg.plugins: diff --git a/autogpt/app.py b/autogpt/app.py index aca97a7512..9eb9d3ab67 100644 --- a/autogpt/app.py +++ b/autogpt/app.py @@ -2,6 +2,7 @@ import json from typing import List, NoReturn, Union from autogpt.agent.agent_manager import AgentManager +from autogpt.commands.command import command, CommandRegistry from autogpt.commands.evaluate_code import evaluate_code from autogpt.commands.google_search import google_official_search, google_search from autogpt.commands.improve_code import improve_code @@ -106,7 +107,7 @@ def map_command_synonyms(command_name: str): return command_name -def execute_command(command_name: str, arguments, prompt: PromptGenerator): +def execute_command(command_registry: CommandRegistry, command_name: str, arguments, prompt: PromptGenerator): """Execute the command and return the result Args: @@ -118,6 +119,13 @@ def execute_command(command_name: str, arguments, prompt: PromptGenerator): memory = get_memory(CFG) try: + cmd = command_registry.commands.get(command_name) + + # If the command is found, call it with the provided arguments + if cmd: + return cmd(**arguments) + + # TODO: Remove commands below after they are moved to the command registry. command_name = map_command_synonyms(command_name) if command_name == "google": # Check if the Google API key is set and use the official search method @@ -248,6 +256,7 @@ def shutdown() -> NoReturn: quit() +@command("start_agent", "Start GPT Agent", '"name": "", "task": "", "prompt": ""') def start_agent(name: str, task: str, prompt: str, model=CFG.fast_llm_model) -> str: """Start an agent with a given name, task, and prompt @@ -280,6 +289,7 @@ def start_agent(name: str, task: str, prompt: str, model=CFG.fast_llm_model) -> return f"Agent {name} created with key {key}. First response: {agent_response}" +@command("message_agent", "Message GPT Agent", '"key": "", "message": ""') def message_agent(key: str, message: str) -> str: """Message an agent with a given key and message""" # Check if the key is a valid integer @@ -294,6 +304,7 @@ def message_agent(key: str, message: str) -> str: return agent_response +@command("list_agents", "List GPT Agents", "") def list_agents(): """List all agents @@ -305,6 +316,7 @@ def list_agents(): ) +@command("delete_agent", "Delete GPT Agent", '"key": ""') def delete_agent(key: str) -> str: """Delete an agent with a given key diff --git a/autogpt/commands/command.py b/autogpt/commands/command.py new file mode 100644 index 0000000000..84a5899008 --- /dev/null +++ b/autogpt/commands/command.py @@ -0,0 +1,123 @@ +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. + 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(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: + 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[name] + + 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: + """ + 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 import_commands(self, module_name: str) -> None: + """ + Imports the specified Python module containing command plugins. + + 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: + module_name (str): The name of the module to import for command plugins. + """ + + 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.""" + 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.command = cmd + + setattr(wrapper, AUTO_GPT_COMMAND_IDENTIFIER, True) + return wrapper + + return decorator diff --git a/autogpt/commands/evaluate_code.py b/autogpt/commands/evaluate_code.py index 8f7cbca9c1..b3d1c87fe2 100644 --- a/autogpt/commands/evaluate_code.py +++ b/autogpt/commands/evaluate_code.py @@ -1,9 +1,11 @@ """Code evaluation module.""" from __future__ import annotations +from autogpt.commands import command from autogpt.llm_utils import call_ai_function +@command("evaluate_code", "Evaluate Code", '"code": ""') def evaluate_code(code: str) -> list[str]: """ A function that takes in a string and returns a response from create chat diff --git a/autogpt/commands/execute_code.py b/autogpt/commands/execute_code.py index eaafa00a82..61d95e36eb 100644 --- a/autogpt/commands/execute_code.py +++ b/autogpt/commands/execute_code.py @@ -4,10 +4,11 @@ import subprocess import docker from docker.errors import ImageNotFound - +from autogpt.commands.command import command from autogpt.workspace import path_in_workspace, WORKSPACE_PATH +@command("execute_python_file", "Execute Python File", '"file": ""') def execute_python_file(file: str): """Execute a Python file in a Docker container and return the output diff --git a/autogpt/commands/file_operations.py b/autogpt/commands/file_operations.py index 60ede701ee..fea62fad18 100644 --- a/autogpt/commands/file_operations.py +++ b/autogpt/commands/file_operations.py @@ -5,6 +5,7 @@ import os import os.path from pathlib import Path from typing import Generator +from autogpt.commands.command import command from autogpt.workspace import path_in_workspace, WORKSPACE_PATH LOG_FILE = "file_logger.txt" @@ -70,6 +71,7 @@ def split_file( start += max_length - overlap +@command("read_file", "Read file", '"file": ""') def read_file(filename: str) -> str: """Read a file and return the contents @@ -122,6 +124,7 @@ def ingest_file( print(f"Error while ingesting file '{filename}': {str(e)}") +@command("write_to_file", "Write to file", '"file": "", "text": ""') def write_to_file(filename: str, text: str) -> str: """Write text to a file @@ -147,6 +150,7 @@ def write_to_file(filename: str, text: str) -> str: return f"Error: {str(e)}" +@command("append_to_file", "Append to file", '"file": "", "text": ""') def append_to_file(filename: str, text: str, shouldLog: bool = True) -> str: """Append text to a file @@ -170,6 +174,7 @@ def append_to_file(filename: str, text: str, shouldLog: bool = True) -> str: return f"Error: {str(e)}" +@command("delete_file", "Delete file", '"file": ""') def delete_file(filename: str) -> str: """Delete a file @@ -190,6 +195,7 @@ def delete_file(filename: str) -> str: return f"Error: {str(e)}" +@command("search_files", "Search Files", '"directory": ""') def search_files(directory: str) -> list[str]: """Search for files in a directory diff --git a/autogpt/commands/image_gen.py b/autogpt/commands/image_gen.py index 6243616ea8..2b62aa354b 100644 --- a/autogpt/commands/image_gen.py +++ b/autogpt/commands/image_gen.py @@ -7,12 +7,14 @@ from base64 import b64decode import openai import requests from PIL import Image +from autogpt.commands.command import command from autogpt.config import Config from autogpt.workspace import path_in_workspace CFG = Config() +@command("generate_image", "Generate Image", '"prompt": ""') def generate_image(prompt: str) -> str: """Generate an image from a prompt. diff --git a/autogpt/commands/improve_code.py b/autogpt/commands/improve_code.py index e3440d8b7c..9d7dc2f0dc 100644 --- a/autogpt/commands/improve_code.py +++ b/autogpt/commands/improve_code.py @@ -2,9 +2,11 @@ from __future__ import annotations import json +from autogpt.commands import command from autogpt.llm_utils import call_ai_function +@command("improve_code", "Get Improved Code", '"suggestions": "", "code": ""') def improve_code(suggestions: list[str], code: str) -> str: """ A function that takes in code and suggestions and returns a response from create diff --git a/autogpt/commands/web_selenium.py b/autogpt/commands/web_selenium.py index 1d078d76d7..a5369ea2d6 100644 --- a/autogpt/commands/web_selenium.py +++ b/autogpt/commands/web_selenium.py @@ -22,6 +22,7 @@ FILE_DIR = Path(__file__).parent.parent CFG = Config() +@command("browse_website", "Browse Website", '"url": "", "question": ""') def browse_website(url: str, question: str) -> tuple[str, WebDriver]: """Browse a website and return the answer and links to the user diff --git a/autogpt/commands/write_tests.py b/autogpt/commands/write_tests.py index 138a1adb6f..f73311787d 100644 --- a/autogpt/commands/write_tests.py +++ b/autogpt/commands/write_tests.py @@ -2,9 +2,11 @@ from __future__ import annotations import json +from autogpt.commands import command from autogpt.llm_utils import call_ai_function +@command("write_tests", "Write Tests", '"code": "", "focus": ""') def write_tests(code: str, focus: list[str]) -> str: """ A function that takes in code and focus topics and returns a response from create diff --git a/scripts/__init__.py b/scripts/__init__.py new file mode 100644 index 0000000000..e69de29bb2