diff --git a/autogpt/__main__.py b/autogpt/__main__.py index 60079f75..5fc9a1ea 100644 --- a/autogpt/__main__.py +++ b/autogpt/__main__.py @@ -47,11 +47,11 @@ def main() -> None: 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') + 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) diff --git a/autogpt/app.py b/autogpt/app.py index 9eb9d3ab..1e782626 100644 --- a/autogpt/app.py +++ b/autogpt/app.py @@ -107,7 +107,12 @@ def map_command_synonyms(command_name: str): return command_name -def execute_command(command_registry: CommandRegistry, 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: @@ -124,7 +129,7 @@ def execute_command(command_registry: CommandRegistry, command_name: str, argume # 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": @@ -191,15 +196,6 @@ def execute_command(command_registry: CommandRegistry, command_name: str, argume return 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 == "execute_shell": - if CFG.execute_local_commands: - return execute_shell(arguments["command_line"]) - else: - return ( - "You are not allowed to run local shell commands. To execute" - " shell commands, EXECUTE_LOCAL_COMMANDS must be set to 'True' " - "in your config. Do not attempt to bypass the restriction." - ) elif command_name == "read_audio_from_file": return read_audio_from_file(arguments["file"]) elif command_name == "generate_image": @@ -256,7 +252,11 @@ def shutdown() -> NoReturn: quit() -@command("start_agent", "Start GPT Agent", '"name": "", "task": "", "prompt": ""') +@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 diff --git a/autogpt/commands/command.py b/autogpt/commands/command.py index 84a58990..d1dfc8fd 100644 --- a/autogpt/commands/command.py +++ b/autogpt/commands/command.py @@ -2,7 +2,7 @@ import os import sys import importlib import inspect -from typing import Callable, Any, List +from typing import Callable, Any, List, Optional # Unique identifier for auto-gpt commands AUTO_GPT_COMMAND_IDENTIFIER = "auto_gpt_command" @@ -17,13 +17,25 @@ class Command: 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): + def __init__( + self, + name: str, + description: str, + method: Callable[..., Any], + signature: str = None, + enabled: bool = True, + disabled_reason: Optional[str] = None, + ): self.name = name self.description = description self.method = method self.signature = signature if signature else str(inspect.signature(self.method)) + self.enabled = enabled + self.disabled_reason = disabled_reason def __call__(self, *args, **kwargs) -> Any: + if not self.enabled: + return f"Command '{self.name}' is disabled: {self.disabled_reason}" return self.method(*args, **kwargs) def __str__(self) -> str: @@ -78,7 +90,9 @@ class CommandRegistry: """ 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())] + 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: @@ -99,18 +113,36 @@ class CommandRegistry: 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): + 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: + 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]: +def command( + name: str, + description: str, + signature: str = None, + enabled: bool = True, + disabled_reason: Optional[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) + cmd = Command( + name=name, + description=description, + method=func, + signature=signature, + enabled=enabled, + disabled_reason=disabled_reason, + ) def wrapper(*args, **kwargs) -> Any: return func(*args, **kwargs) diff --git a/autogpt/commands/evaluate_code.py b/autogpt/commands/evaluate_code.py index b3d1c87f..1c9b117d 100644 --- a/autogpt/commands/evaluate_code.py +++ b/autogpt/commands/evaluate_code.py @@ -1,7 +1,7 @@ """Code evaluation module.""" from __future__ import annotations -from autogpt.commands import command +from autogpt.commands.command import command from autogpt.llm_utils import call_ai_function diff --git a/autogpt/commands/execute_code.py b/autogpt/commands/execute_code.py index 61d95e36..aa8a3545 100644 --- a/autogpt/commands/execute_code.py +++ b/autogpt/commands/execute_code.py @@ -4,9 +4,12 @@ import subprocess import docker from docker.errors import ImageNotFound +from autogpt.config import Config from autogpt.commands.command import command from autogpt.workspace import path_in_workspace, WORKSPACE_PATH +CFG = Config() + @command("execute_python_file", "Execute Python File", '"file": ""') def execute_python_file(file: str): @@ -89,6 +92,15 @@ def execute_python_file(file: str): return f"Error: {str(e)}" +@command( + "execute_shell", + "Execute Shell Command, non-interactive commands only", + '"command_line": ""', + CFG.execute_local_commands, + "You are not allowed to run local shell commands. To execute" + " shell commands, EXECUTE_LOCAL_COMMANDS must be set to 'True' " + "in your config. Do not attempt to bypass the restriction.", +) def execute_shell(command_line: str) -> str: """Execute a shell command and return the output @@ -98,6 +110,13 @@ def execute_shell(command_line: str) -> str: Returns: str: The output of the command """ + + if not CFG.execute_local_commands: + return ( + "You are not allowed to run local shell commands. To execute" + " shell commands, EXECUTE_LOCAL_COMMANDS must be set to 'True' " + "in your config. Do not attempt to bypass the restriction." + ) current_dir = os.getcwd() # Change dir into workspace if necessary if str(WORKSPACE_PATH) not in current_dir: diff --git a/autogpt/commands/improve_code.py b/autogpt/commands/improve_code.py index 9d7dc2f0..0bfe7253 100644 --- a/autogpt/commands/improve_code.py +++ b/autogpt/commands/improve_code.py @@ -6,7 +6,11 @@ from autogpt.commands import command from autogpt.llm_utils import call_ai_function -@command("improve_code", "Get Improved Code", '"suggestions": "", "code": ""') +@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 a5369ea2..591d3162 100644 --- a/autogpt/commands/web_selenium.py +++ b/autogpt/commands/web_selenium.py @@ -22,7 +22,11 @@ FILE_DIR = Path(__file__).parent.parent CFG = Config() -@command("browse_website", "Browse Website", '"url": "", "question": ""') +@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 f7331178..23d4c130 100644 --- a/autogpt/commands/write_tests.py +++ b/autogpt/commands/write_tests.py @@ -6,7 +6,11 @@ from autogpt.commands import command from autogpt.llm_utils import call_ai_function -@command("write_tests", "Write Tests", '"code": "", "focus": ""') +@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/autogpt/prompts/prompt.py b/autogpt/prompts/prompt.py index 8dacaf7f..d82cdb16 100644 --- a/autogpt/prompts/prompt.py +++ b/autogpt/prompts/prompt.py @@ -87,16 +87,6 @@ def build_default_prompt_generator() -> PromptGenerator: ("Convert Audio to text", "read_audio_from_file", {"file": ""}), ) - # Only add shell command to the prompt if the AI is allowed to execute it - if CFG.execute_local_commands: - commands.append( - ( - "Execute Shell Command, non-interactive commands only", - "execute_shell", - {"command_line": ""}, - ), - ) - # Add these command last. commands.append( ("Do Nothing", "do_nothing", {}), diff --git a/tests/mocks/mock_commands.py b/tests/mocks/mock_commands.py index d68ceb81..d64284bc 100644 --- a/tests/mocks/mock_commands.py +++ b/tests/mocks/mock_commands.py @@ -1,6 +1,6 @@ -from auto_gpt.commands import Command, command +from autogpt.commands.command import command -@command('function_based', 'Function-based test command') +@command("function_based", "Function-based test command") def function_based(arg1: int, arg2: str) -> str: - return f'{arg1} - {arg2}' + return f"{arg1} - {arg2}" diff --git a/tests/test_commands.py b/tests/test_commands.py index 9ff0cd4c..a21bbb4d 100644 --- a/tests/test_commands.py +++ b/tests/test_commands.py @@ -3,7 +3,7 @@ import sys from pathlib import Path import pytest -from auto_gpt.commands import Command, CommandRegistry +from autogpt.commands.command import Command, CommandRegistry class TestCommand: @@ -12,7 +12,9 @@ class TestCommand: return f"{arg1} - {arg2}" def test_command_creation(self): - cmd = Command(name="example", description="Example command", method=self.example_function) + cmd = Command( + name="example", description="Example command", method=self.example_function + ) assert cmd.name == "example" assert cmd.description == "Example command" @@ -20,30 +22,40 @@ class TestCommand: assert cmd.signature == "(arg1: int, arg2: str) -> str" def test_command_call(self): - cmd = Command(name="example", description="Example command", method=self.example_function) + 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) + 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) + 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) + 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: @@ -52,7 +64,9 @@ class TestCommandRegistry: 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) + cmd = Command( + name="example", description="Example command", method=self.example_function + ) registry.register(cmd) @@ -62,7 +76,9 @@ class TestCommandRegistry: 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) + cmd = Command( + name="example", description="Example command", method=self.example_function + ) registry.register(cmd) registry.unregister(cmd.name) @@ -72,7 +88,9 @@ class TestCommandRegistry: 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) + cmd = Command( + name="example", description="Example command", method=self.example_function + ) registry.register(cmd) retrieved_cmd = registry.get_command(cmd.name) @@ -89,7 +107,9 @@ class TestCommandRegistry: 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) + cmd = Command( + name="example", description="Example command", method=self.example_function + ) registry.register(cmd) result = registry.call("example", arg1=1, arg2="test") @@ -106,7 +126,9 @@ class TestCommandRegistry: 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) + cmd = Command( + name="example", description="Example command", method=self.example_function + ) registry.register(cmd) command_prompt = registry.command_prompt() @@ -122,7 +144,10 @@ class TestCommandRegistry: assert "function_based" in registry.commands assert registry.commands["function_based"].name == "function_based" - assert registry.commands["function_based"].description == "Function-based test command" + 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.""" @@ -144,4 +169,7 @@ class TestCommandRegistry: assert "function_based" in registry.commands assert registry.commands["function_based"].name == "function_based" - assert registry.commands["function_based"].description == "Function-based test command" + assert ( + registry.commands["function_based"].description + == "Function-based test command" + )