Add fields to disable the command if needed by configuration, blacked.

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.ai_functions")
-    command_registry.import_commands('scripts.commands')
+    command_registry.import_commands("scripts.commands")
-    command_registry.import_commands('scripts.execute_code')
+    command_registry.import_commands("scripts.execute_code")
-    command_registry.import_commands('scripts.agent_manager')
+    command_registry.import_commands("scripts.agent_manager")
-    command_registry.import_commands('scripts.file_operations')
+    command_registry.import_commands("scripts.file_operations")
     ai_name = ""
     ai_config = construct_main_ai_config()
     # print(prompt)

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": "<name>", "task": "<short_task_desc>", "prompt": "<prompt>"')
+@command(
+    "start_agent",
+    "Start GPT Agent",
+    '"name": "<name>", "task": "<short_task_desc>", "prompt": "<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
 

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)

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
 
 

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": "<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": "<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:

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": "<list_of_suggestions>", "code": "<full_code_string>"')
+@command(
+    "improve_code",
+    "Get Improved Code",
+    '"suggestions": "<list_of_suggestions>", "code": "<full_code_string>"',
+)
 def improve_code(suggestions: list[str], code: str) -> str:
     """
     A function that takes in code and suggestions and returns a response from create

autogpt/commands/web_selenium.py

@@ -22,7 +22,11 @@ FILE_DIR = Path(__file__).parent.parent
 CFG = Config()
 
 
-@command("browse_website", "Browse Website", '"url": "<url>", "question": "<what_you_want_to_find_on_website>"')
+@command(
+    "browse_website",
+    "Browse Website",
+    '"url": "<url>", "question": "<what_you_want_to_find_on_website>"',
+)
 def browse_website(url: str, question: str) -> tuple[str, WebDriver]:
     """Browse a website and return the answer and links to the user
 

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": "<full_code_string>", "focus": "<list_of_focus_areas>"')
+@command(
+    "write_tests",
+    "Write Tests",
+    '"code": "<full_code_string>", "focus": "<list_of_focus_areas>"',
+)
 def write_tests(code: str, focus: list[str]) -> str:
     """
     A function that takes in code and focus topics and returns a response from create

autogpt/prompts/prompt.py

@@ -87,16 +87,6 @@ def build_default_prompt_generator() -> PromptGenerator:
             ("Convert Audio to text", "read_audio_from_file", {"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": "<command_line>"},
-            ),
-        )
-
     # Add these command last.
     commands.append(
         ("Do Nothing", "do_nothing", {}),

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}"

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"
+        )