From 6470e44bd9c13014aab0088fa6ce2e692d9c31e2 Mon Sep 17 00:00:00 2001
From: Erik Peterson <e@eriklp.com>
Date: Tue, 13 Jun 2023 19:10:10 -0700
Subject: [PATCH] Refactor module layout of command classes

---
 autogpt/agent/agent.py                        |  2 +-
 autogpt/app.py                                |  3 +-
 autogpt/command_decorator.py                  | 51 +++++++++++
 autogpt/commands/analyze_code.py              |  2 +-
 autogpt/commands/audio_text.py                |  2 +-
 autogpt/commands/execute_code.py              |  2 +-
 autogpt/commands/file_operations.py           |  2 +-
 autogpt/commands/git_operations.py            |  2 +-
 autogpt/commands/google_search.py             |  2 +-
 autogpt/commands/image_gen.py                 |  2 +-
 autogpt/commands/improve_code.py              |  2 +-
 autogpt/commands/task_statuses.py             |  2 +-
 autogpt/commands/web_selenium.py              |  2 +-
 autogpt/commands/write_tests.py               |  2 +-
 autogpt/config/ai_config.py                   |  2 +-
 autogpt/main.py                               |  2 +-
 autogpt/models/__init__.py                    |  0
 autogpt/models/command.py                     | 41 +++++++++
 .../command.py => models/command_registry.py} | 89 +------------------
 autogpt/prompts/generator.py                  |  2 +-
 tests/conftest.py                             |  2 +-
 tests/integration/agent_factory.py            |  2 +-
 tests/mocks/mock_commands.py                  |  2 +-
 tests/unit/test_commands.py                   |  3 +-
 24 files changed, 117 insertions(+), 106 deletions(-)
 create mode 100644 autogpt/command_decorator.py
 create mode 100644 autogpt/models/__init__.py
 create mode 100644 autogpt/models/command.py
 rename autogpt/{commands/command.py => models/command_registry.py} (56%)

diff --git a/autogpt/agent/agent.py b/autogpt/agent/agent.py
index 5a236f673e1..2fed0d4bc65 100644
--- a/autogpt/agent/agent.py
+++ b/autogpt/agent/agent.py
@@ -5,7 +5,6 @@
 
 from colorama import Fore, Style
 
-from autogpt.commands.command import CommandRegistry
 from autogpt.config import Config
 from autogpt.config.ai_config import AIConfig
 from autogpt.json_utils.utilities import extract_json_from_response, validate_json
@@ -24,6 +23,7 @@
 from autogpt.logs import logger, print_assistant_thoughts
 from autogpt.memory.message_history import MessageHistory
 from autogpt.memory.vector import VectorMemory
+from autogpt.models.command_registry import CommandRegistry
 from autogpt.speech import say_text
 from autogpt.spinner import Spinner
 from autogpt.utils import clean_input
diff --git a/autogpt/app.py b/autogpt/app.py
index c081835cf64..fee3413a17b 100644
--- a/autogpt/app.py
+++ b/autogpt/app.py
@@ -4,8 +4,9 @@
 
 from autogpt.agent.agent import Agent
 from autogpt.agent.agent_manager import AgentManager
-from autogpt.commands.command import CommandRegistry, command
+from autogpt.command_decorator import command
 from autogpt.commands.web_requests import scrape_links, scrape_text
+from autogpt.models.command_registry import CommandRegistry
 from autogpt.processing.text import summarize_text
 from autogpt.speech import say_text
 from autogpt.url_utils.validators import validate_url
diff --git a/autogpt/command_decorator.py b/autogpt/command_decorator.py
new file mode 100644
index 00000000000..3f8279e4d6e
--- /dev/null
+++ b/autogpt/command_decorator.py
@@ -0,0 +1,51 @@
+import functools
+from typing import Any, Callable, Optional
+
+from autogpt.config import Config
+from autogpt.logs import logger
+from autogpt.models.command import Command
+
+# Unique identifier for auto-gpt commands
+AUTO_GPT_COMMAND_IDENTIFIER = "auto_gpt_command"
+
+
+def command(
+    name: str,
+    description: str,
+    signature: str,
+    enabled: bool | Callable[[Config], bool] = True,
+    disabled_reason: Optional[str] = None,
+) -> Callable[..., Any]:
+    """The command decorator is used to create Command objects from ordinary functions."""
+
+    # TODO: Remove this in favor of better command management
+    CFG = Config()
+
+    if callable(enabled):
+        enabled = enabled(CFG)
+    if not enabled:
+        if disabled_reason is not None:
+            logger.debug(f"Command '{name}' is disabled: {disabled_reason}")
+        return lambda func: func
+
+    def decorator(func: Callable[..., Any]) -> Command:
+        cmd = Command(
+            name=name,
+            description=description,
+            method=func,
+            signature=signature,
+            enabled=enabled,
+            disabled_reason=disabled_reason,
+        )
+
+        @functools.wraps(func)
+        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/analyze_code.py b/autogpt/commands/analyze_code.py
index ca7fcb015f4..cd176a9aa56 100644
--- a/autogpt/commands/analyze_code.py
+++ b/autogpt/commands/analyze_code.py
@@ -2,7 +2,7 @@
 from __future__ import annotations
 
 from autogpt.agent.agent import Agent
-from autogpt.commands.command import command
+from autogpt.command_decorator import command
 from autogpt.llm.utils import call_ai_function
 
 
diff --git a/autogpt/commands/audio_text.py b/autogpt/commands/audio_text.py
index 2991fff32c3..e77e37cc35e 100644
--- a/autogpt/commands/audio_text.py
+++ b/autogpt/commands/audio_text.py
@@ -4,7 +4,7 @@
 import requests
 
 from autogpt.agent.agent import Agent
-from autogpt.commands.command import command
+from autogpt.command_decorator import command
 
 
 @command(
diff --git a/autogpt/commands/execute_code.py b/autogpt/commands/execute_code.py
index 109caa3aa60..c422d652f9d 100644
--- a/autogpt/commands/execute_code.py
+++ b/autogpt/commands/execute_code.py
@@ -7,7 +7,7 @@
 from docker.errors import ImageNotFound
 
 from autogpt.agent.agent import Agent
-from autogpt.commands.command import command
+from autogpt.command_decorator import command
 from autogpt.config import Config
 from autogpt.logs import logger
 from autogpt.setup import CFG
diff --git a/autogpt/commands/file_operations.py b/autogpt/commands/file_operations.py
index 2ff1484474f..d74fee9610c 100644
--- a/autogpt/commands/file_operations.py
+++ b/autogpt/commands/file_operations.py
@@ -13,7 +13,7 @@
 from requests.adapters import HTTPAdapter, Retry
 
 from autogpt.agent.agent import Agent
-from autogpt.commands.command import command
+from autogpt.command_decorator import command
 from autogpt.commands.file_operations_utils import read_textual_file
 from autogpt.logs import logger
 from autogpt.memory.vector import MemoryItem, VectorMemory
diff --git a/autogpt/commands/git_operations.py b/autogpt/commands/git_operations.py
index e844fd4151f..8dfe213cdf8 100644
--- a/autogpt/commands/git_operations.py
+++ b/autogpt/commands/git_operations.py
@@ -3,7 +3,7 @@
 from git.repo import Repo
 
 from autogpt.agent.agent import Agent
-from autogpt.commands.command import command
+from autogpt.command_decorator import command
 from autogpt.url_utils.validators import validate_url
 
 
diff --git a/autogpt/commands/google_search.py b/autogpt/commands/google_search.py
index b9d243f9788..e6a1fc05881 100644
--- a/autogpt/commands/google_search.py
+++ b/autogpt/commands/google_search.py
@@ -8,7 +8,7 @@
 from duckduckgo_search import DDGS
 
 from autogpt.agent.agent import Agent
-from autogpt.commands.command import command
+from autogpt.command_decorator import command
 
 DUCKDUCKGO_MAX_ATTEMPTS = 3
 
diff --git a/autogpt/commands/image_gen.py b/autogpt/commands/image_gen.py
index b2dc9ea4802..5bed8e00bb6 100644
--- a/autogpt/commands/image_gen.py
+++ b/autogpt/commands/image_gen.py
@@ -10,7 +10,7 @@
 from PIL import Image
 
 from autogpt.agent.agent import Agent
-from autogpt.commands.command import command
+from autogpt.command_decorator import command
 from autogpt.logs import logger
 
 
diff --git a/autogpt/commands/improve_code.py b/autogpt/commands/improve_code.py
index 05e9b51c1c0..d4f87782ce7 100644
--- a/autogpt/commands/improve_code.py
+++ b/autogpt/commands/improve_code.py
@@ -3,7 +3,7 @@
 import json
 
 from autogpt.agent.agent import Agent
-from autogpt.commands.command import command
+from autogpt.command_decorator import command
 from autogpt.llm.utils import call_ai_function
 
 
diff --git a/autogpt/commands/task_statuses.py b/autogpt/commands/task_statuses.py
index 283328a3661..d5718fd3595 100644
--- a/autogpt/commands/task_statuses.py
+++ b/autogpt/commands/task_statuses.py
@@ -4,7 +4,7 @@
 from typing import NoReturn
 
 from autogpt.agent.agent import Agent
-from autogpt.commands.command import command
+from autogpt.command_decorator import command
 from autogpt.logs import logger
 
 
diff --git a/autogpt/commands/web_selenium.py b/autogpt/commands/web_selenium.py
index 14036c85e09..bdc5e613753 100644
--- a/autogpt/commands/web_selenium.py
+++ b/autogpt/commands/web_selenium.py
@@ -28,7 +28,7 @@
 from webdriver_manager.microsoft import EdgeChromiumDriverManager as EdgeDriverManager
 
 from autogpt.agent.agent import Agent
-from autogpt.commands.command import command
+from autogpt.command_decorator import command
 from autogpt.logs import logger
 from autogpt.memory.vector import MemoryItem, get_memory
 from autogpt.processing.html import extract_hyperlinks, format_hyperlinks
diff --git a/autogpt/commands/write_tests.py b/autogpt/commands/write_tests.py
index c09930b9f3c..881b6ac4f1f 100644
--- a/autogpt/commands/write_tests.py
+++ b/autogpt/commands/write_tests.py
@@ -4,7 +4,7 @@
 import json
 
 from autogpt.agent.agent import Agent
-from autogpt.commands.command import command
+from autogpt.command_decorator import command
 from autogpt.llm.utils import call_ai_function
 
 
diff --git a/autogpt/config/ai_config.py b/autogpt/config/ai_config.py
index 1a526832349..d118be3f419 100644
--- a/autogpt/config/ai_config.py
+++ b/autogpt/config/ai_config.py
@@ -13,7 +13,7 @@
 import yaml
 
 if TYPE_CHECKING:
-    from autogpt.commands.command import CommandRegistry
+    from autogpt.models.command_registry import CommandRegistry
     from autogpt.prompts.generator import PromptGenerator
 
 # Soon this will go in a folder where it remembers more stuff about the run(s)
diff --git a/autogpt/main.py b/autogpt/main.py
index ce6a983dc4f..97baa7b28b7 100644
--- a/autogpt/main.py
+++ b/autogpt/main.py
@@ -6,11 +6,11 @@
 from colorama import Fore, Style
 
 from autogpt.agent import Agent
-from autogpt.commands.command import CommandRegistry
 from autogpt.config import Config, check_openai_api_key
 from autogpt.configurator import create_config
 from autogpt.logs import logger
 from autogpt.memory.vector import get_memory
+from autogpt.models.command_registry import CommandRegistry
 from autogpt.plugins import scan_plugins
 from autogpt.prompts.prompt import DEFAULT_TRIGGERING_PROMPT, construct_main_ai_config
 from autogpt.utils import (
diff --git a/autogpt/models/__init__.py b/autogpt/models/__init__.py
new file mode 100644
index 00000000000..e69de29bb2d
diff --git a/autogpt/models/command.py b/autogpt/models/command.py
new file mode 100644
index 00000000000..a925ca04362
--- /dev/null
+++ b/autogpt/models/command.py
@@ -0,0 +1,41 @@
+from typing import Any, Callable, Optional
+
+from autogpt.config import Config
+
+
+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 = "",
+        enabled: bool | Callable[[Config], bool] = True,
+        disabled_reason: Optional[str] = None,
+    ):
+        self.name = name
+        self.description = description
+        self.method = method
+        self.signature = signature
+        self.enabled = enabled
+        self.disabled_reason = disabled_reason
+
+    def __call__(self, *args, **kwargs) -> Any:
+        if hasattr(kwargs, "config") and callable(self.enabled):
+            self.enabled = self.enabled(kwargs["config"])
+        if not self.enabled:
+            if self.disabled_reason:
+                return f"Command '{self.name}' is disabled: {self.disabled_reason}"
+            return f"Command '{self.name}' is disabled"
+        return self.method(*args, **kwargs)
+
+    def __str__(self) -> str:
+        return f"{self.name}: {self.description}, args: {self.signature}"
diff --git a/autogpt/commands/command.py b/autogpt/models/command_registry.py
similarity index 56%
rename from autogpt/commands/command.py
rename to autogpt/models/command_registry.py
index 742cc8df649..29d0143d91a 100644
--- a/autogpt/commands/command.py
+++ b/autogpt/models/command_registry.py
@@ -1,51 +1,10 @@
-import functools
 import importlib
 import inspect
-from typing import Any, Callable, Optional
+from typing import Any, Callable
 
-from autogpt.config import Config
+from autogpt.command_decorator import AUTO_GPT_COMMAND_IDENTIFIER
 from autogpt.logs import logger
-
-# 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 = "",
-        enabled: bool | Callable[[Config], bool] = True,
-        disabled_reason: Optional[str] = None,
-    ):
-        self.name = name
-        self.description = description
-        self.method = method
-        self.signature = signature
-        self.enabled = enabled
-        self.disabled_reason = disabled_reason
-
-    def __call__(self, *args, **kwargs) -> Any:
-        if hasattr(kwargs, "config") and callable(self.enabled):
-            self.enabled = self.enabled(kwargs["config"])
-        if not self.enabled:
-            if self.disabled_reason:
-                return f"Command '{self.name}' is disabled: {self.disabled_reason}"
-            return f"Command '{self.name}' is disabled"
-        return self.method(*args, **kwargs)
-
-    def __str__(self) -> str:
-        return f"{self.name}: {self.description}, args: {self.signature}"
+from autogpt.models.command import Command
 
 
 class CommandRegistry:
@@ -133,45 +92,3 @@ def import_commands(self, module_name: str) -> None:
             ):
                 cmd_instance = attr()
                 self.register(cmd_instance)
-
-
-def command(
-    name: str,
-    description: str,
-    signature: str,
-    enabled: bool | Callable[[Config], bool] = True,
-    disabled_reason: Optional[str] = None,
-) -> Callable[..., Any]:
-    """The command decorator is used to create Command objects from ordinary functions."""
-
-    # TODO: Remove this in favor of better command management
-    CFG = Config()
-
-    if callable(enabled):
-        enabled = enabled(CFG)
-    if not enabled:
-        if disabled_reason is not None:
-            logger.debug(f"Command '{name}' is disabled: {disabled_reason}")
-        return lambda func: func
-
-    def decorator(func: Callable[..., Any]) -> Command:
-        cmd = Command(
-            name=name,
-            description=description,
-            method=func,
-            signature=signature,
-            enabled=enabled,
-            disabled_reason=disabled_reason,
-        )
-
-        @functools.wraps(func)
-        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/prompts/generator.py b/autogpt/prompts/generator.py
index 7101acfea11..2a0334bf45a 100644
--- a/autogpt/prompts/generator.py
+++ b/autogpt/prompts/generator.py
@@ -4,7 +4,7 @@
 from autogpt.json_utils.utilities import llm_response_schema
 
 if TYPE_CHECKING:
-    from autogpt.commands.command import CommandRegistry
+    from autogpt.models.command_registry import CommandRegistry
 
 
 class PromptGenerator:
diff --git a/tests/conftest.py b/tests/conftest.py
index 671096fdc24..97620e2173e 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -7,12 +7,12 @@
 from pytest_mock import MockerFixture
 
 from autogpt.agent.agent import Agent
-from autogpt.commands.command import CommandRegistry
 from autogpt.config.ai_config import AIConfig
 from autogpt.config.config import Config
 from autogpt.llm.api_manager import ApiManager
 from autogpt.logs import TypingConsoleHandler
 from autogpt.memory.vector import get_memory
+from autogpt.models.command_registry import CommandRegistry
 from autogpt.prompts.prompt import DEFAULT_TRIGGERING_PROMPT
 from autogpt.workspace import Workspace
 
diff --git a/tests/integration/agent_factory.py b/tests/integration/agent_factory.py
index c9b99f50f71..9078a8434da 100644
--- a/tests/integration/agent_factory.py
+++ b/tests/integration/agent_factory.py
@@ -1,10 +1,10 @@
 import pytest
 
 from autogpt.agent import Agent
-from autogpt.commands.command import CommandRegistry
 from autogpt.config import AIConfig, Config
 from autogpt.main import COMMAND_CATEGORIES
 from autogpt.memory.vector import NoMemory, get_memory
+from autogpt.models.command_registry import CommandRegistry
 from autogpt.prompts.prompt import DEFAULT_TRIGGERING_PROMPT
 from autogpt.workspace import Workspace
 
diff --git a/tests/mocks/mock_commands.py b/tests/mocks/mock_commands.py
index 42b0ea1134c..7b16f1d1502 100644
--- a/tests/mocks/mock_commands.py
+++ b/tests/mocks/mock_commands.py
@@ -1,4 +1,4 @@
-from autogpt.commands.command import command
+from autogpt.command_decorator import command
 
 
 @command(
diff --git a/tests/unit/test_commands.py b/tests/unit/test_commands.py
index 5779a8a3f00..f02cb620e43 100644
--- a/tests/unit/test_commands.py
+++ b/tests/unit/test_commands.py
@@ -5,7 +5,8 @@
 
 import pytest
 
-from autogpt.commands.command import Command, CommandRegistry
+from autogpt.models.command import Command
+from autogpt.models.command_registry import CommandRegistry
 
 SIGNATURE = "(arg1: int, arg2: str) -> str"