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

This commit is contained in:
BillSchumacher
2023-04-16 15:49:36 -05:00
parent df5cc3303f
commit 167628c696
11 changed files with 136 additions and 55 deletions

View File

@@ -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)

View File

@@ -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

View File

@@ -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)

View File

@@ -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

View File

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

View File

@@ -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

View File

@@ -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

View File

@@ -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

View File

@@ -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", {}),

View File

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

View File

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