Add categories to command registry (#5063)

* Add categories to command registry

* Fix tests

agbenchmark/benchmarks.py

@@ -50,8 +50,8 @@ def get_command_registry(config: Config):
     enabled_command_categories = [
         x for x in COMMAND_CATEGORIES if x not in config.disabled_command_categories
     ]
-    for command_category in enabled_command_categories:
-        command_registry.import_commands(command_category)
+    for command_module in enabled_command_categories:
+        command_registry.import_command_module(command_module)
     return command_registry
 
 

autogpt/app/main.py

@@ -134,36 +134,9 @@ def run_auto_gpt(
     config.file_logger_path = Workspace.build_file_logger_path(config.workspace_path)
 
     config.plugins = scan_plugins(config, config.debug_mode)
+
     # Create a CommandRegistry instance and scan default folder
-    command_registry = CommandRegistry()
-
-    logger.debug(
-        f"The following command categories are disabled: {config.disabled_command_categories}"
-    )
-    enabled_command_categories = [
-        x for x in COMMAND_CATEGORIES if x not in config.disabled_command_categories
-    ]
-
-    logger.debug(
-        f"The following command categories are enabled: {enabled_command_categories}"
-    )
-
-    for command_category in enabled_command_categories:
-        command_registry.import_commands(command_category)
-
-    # Unregister commands that are incompatible with the current config
-    incompatible_commands = []
-    for command in command_registry.commands.values():
-        if callable(command.enabled) and not command.enabled(config):
-            command.enabled = False
-            incompatible_commands.append(command)
-
-    for command in incompatible_commands:
-        command_registry.unregister(command)
-        logger.debug(
-            f"Unregistering incompatible command: {command.name}, "
-            f"reason - {command.disabled_reason or 'Disabled by current config.'}"
-        )
+    command_registry = CommandRegistry.with_command_modules(COMMAND_CATEGORIES, config)
 
     ai_config = construct_main_ai_config(
         config,

autogpt/commands/__init__.py

@@ -3,5 +3,5 @@ COMMAND_CATEGORIES = [
     "autogpt.commands.file_operations",
     "autogpt.commands.web_search",
     "autogpt.commands.web_selenium",
-    "autogpt.commands.task_statuses",
+    "autogpt.commands.system",
 ]

autogpt/commands/execute_code.py

@@ -1,4 +1,8 @@
-"""Execute code in a Docker container"""
+"""Commands to execute code"""
+
+COMMAND_CATEGORY = "execute_code"
+COMMAND_CATEGORY_TITLE = "Execute Code"
+
 import os
 import subprocess
 from pathlib import Path

autogpt/commands/file_operations.py

@@ -1,6 +1,10 @@
-"""File operations for AutoGPT"""
+"""Commands to perform operations on files"""
+
 from __future__ import annotations
 
+COMMAND_CATEGORY = "file_operations"
+COMMAND_CATEGORY_TITLE = "File Operations"
+
 import contextlib
 import hashlib
 import os

autogpt/commands/git_operations.py

@@ -1,4 +1,7 @@
-"""Git operations for autogpt"""
+"""Commands to perform Git operations"""
+
+COMMAND_CATEGORY = "git_operations"
+COMMAND_CATEGORY_TITLE = "Git Operations"
 
 from git.repo import Repo
 

autogpt/commands/image_gen.py

@@ -1,4 +1,8 @@
-""" Image Generation Module for AutoGPT."""
+"""Commands to generate images based on text input"""
+
+COMMAND_CATEGORY = "text_to_image"
+COMMAND_CATEGORY_TITLE = "Text to Image"
+
 import io
 import json
 import time

autogpt/commands/system.py

@@ -1,6 +1,10 @@
-"""Task Statuses module."""
+"""Commands to control the internal state of the program"""
+
 from __future__ import annotations
 
+COMMAND_CATEGORY = "system"
+COMMAND_CATEGORY_TITLE = "System"
+
 from typing import NoReturn
 
 from autogpt.agents.agent import Agent

autogpt/commands/web_search.py

@@ -1,6 +1,10 @@
-"""Google search command for Autogpt."""
+"""Commands to search the web with"""
+
 from __future__ import annotations
 
+COMMAND_CATEGORY = "web_search"
+COMMAND_CATEGORY_TITLE = "Web Search"
+
 import json
 import time
 from itertools import islice

autogpt/commands/web_selenium.py

@@ -1,6 +1,10 @@
-"""Selenium web scraping module."""
+"""Commands for browsing a website"""
+
 from __future__ import annotations
 
+COMMAND_CATEGORY = "web_browse"
+COMMAND_CATEGORY_TITLE = "Web Browsing"
+
 import logging
 from pathlib import Path
 from sys import platform

autogpt/core/runner/cli_web_app/server/api.py

@@ -109,6 +109,6 @@ def get_command_registry(config: Config):
     enabled_command_categories = [
         x for x in COMMAND_CATEGORIES if x not in config.disabled_command_categories
     ]
-    for command_category in enabled_command_categories:
-        command_registry.import_commands(command_category)
+    for command_module in enabled_command_categories:
+        command_registry.import_command_module(command_module)
     return command_registry

autogpt/models/command_registry.py

@@ -1,6 +1,13 @@
+from __future__ import annotations
+
 import importlib
 import inspect
-from typing import Any
+from dataclasses import dataclass, field
+from types import ModuleType
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from autogpt.config import Config
 
 from autogpt.command_decorator import AUTO_GPT_COMMAND_IDENTIFIER
 from autogpt.logs import logger
@@ -18,9 +25,21 @@ class CommandRegistry:
     commands: dict[str, Command]
     commands_aliases: dict[str, Command]
 
+    # Alternative way to structure the registry; currently redundant with self.commands
+    categories: dict[str, CommandCategory]
+
+    @dataclass
+    class CommandCategory:
+        name: str
+        title: str
+        description: str
+        commands: list[Command] = field(default_factory=list[Command])
+        modules: list[ModuleType] = field(default_factory=list[ModuleType])
+
     def __init__(self):
         self.commands = {}
         self.commands_aliases = {}
+        self.categories = {}
 
     def __contains__(self, command_name: str):
         return command_name in self.commands or command_name in self.commands_aliases
@@ -84,7 +103,41 @@ class CommandRegistry:
         ]
         return "\n".join(commands_list)
 
-    def import_commands(self, module_name: str) -> None:
+    @staticmethod
+    def with_command_modules(modules: list[str], config: Config) -> CommandRegistry:
+        new_registry = CommandRegistry()
+
+        logger.debug(
+            f"The following command categories are disabled: {config.disabled_command_categories}"
+        )
+        enabled_command_modules = [
+            x for x in modules if x not in config.disabled_command_categories
+        ]
+
+        logger.debug(
+            f"The following command categories are enabled: {enabled_command_modules}"
+        )
+
+        for command_module in enabled_command_modules:
+            new_registry.import_command_module(command_module)
+
+        # Unregister commands that are incompatible with the current config
+        incompatible_commands: list[Command] = []
+        for command in new_registry.commands.values():
+            if callable(command.enabled) and not command.enabled(config):
+                command.enabled = False
+                incompatible_commands.append(command)
+
+        for command in incompatible_commands:
+            new_registry.unregister(command)
+            logger.debug(
+                f"Unregistering incompatible command: {command.name}, "
+                f"reason - {command.disabled_reason or 'Disabled by current config.'}"
+            )
+
+        return new_registry
+
+    def import_command_module(self, module_name: str) -> None:
         """
         Imports the specified Python module containing command plugins.
 
@@ -99,16 +152,42 @@ class CommandRegistry:
 
         module = importlib.import_module(module_name)
 
+        category = self.register_module_category(module)
+
         for attr_name in dir(module):
             attr = getattr(module, attr_name)
+
+            command = None
+
             # Register decorated functions
-            if hasattr(attr, AUTO_GPT_COMMAND_IDENTIFIER) and getattr(
-                attr, AUTO_GPT_COMMAND_IDENTIFIER
-            ):
-                self.register(attr.command)
+            if getattr(attr, AUTO_GPT_COMMAND_IDENTIFIER, False):
+                command = attr.command
+
             # Register command classes
             elif (
                 inspect.isclass(attr) and issubclass(attr, Command) and attr != Command
             ):
-                cmd_instance = attr()
-                self.register(cmd_instance)
+                command = attr()
+
+            if command:
+                self.register(command)
+                category.commands.append(command)
+
+    def register_module_category(self, module: ModuleType) -> CommandCategory:
+        if not (category_name := getattr(module, "COMMAND_CATEGORY", None)):
+            raise ValueError(f"Cannot import invalid command module {module.__name__}")
+
+        if category_name not in self.categories:
+            self.categories[category_name] = CommandRegistry.CommandCategory(
+                name=category_name,
+                title=getattr(
+                    module, "COMMAND_CATEGORY_TITLE", category_name.capitalize()
+                ),
+                description=getattr(module, "__doc__", ""),
+            )
+
+        category = self.categories[category_name]
+        if module not in category.modules:
+            category.modules.append(module)
+
+        return category

tests/mocks/mock_commands.py

@@ -1,5 +1,7 @@
 from autogpt.command_decorator import command
 
+COMMAND_CATEGORY = "mock"
+
 
 @command(
     "function_based",

tests/unit/test_commands.py

@@ -193,7 +193,7 @@ def test_import_mock_commands_module():
     registry = CommandRegistry()
     mock_commands_module = "tests.mocks.mock_commands"
 
-    registry.import_commands(mock_commands_module)
+    registry.import_command_module(mock_commands_module)
 
     assert "function_based" in registry
     assert registry.commands["function_based"].name == "function_based"
@@ -219,7 +219,7 @@ def test_import_temp_command_file_module(tmp_path: Path):
     sys.path.append(str(tmp_path))
 
     temp_commands_module = "mock_commands"
-    registry.import_commands(temp_commands_module)
+    registry.import_command_module(temp_commands_module)
 
     # Remove the temp directory from sys.path
     sys.path.remove(str(tmp_path))