Added ability to run addon ChatCommands

esbmc_ai/__main__.py

@@ -3,7 +3,6 @@
 # Author: Yiannis Charalambous 2023
 
 from pathlib import Path
-import re
 import sys
 
 # Enables arrow key functionality for input(). Do not remove import.
@@ -13,6 +12,10 @@
 
 from langchain_core.language_models import BaseChatModel
 
+from esbmc_ai.command_runner import CommandRunner
+from esbmc_ai.commands.fix_code_command import FixCodeCommandResult
+
+
 import argparse
 
 from esbmc_ai import Config
@@ -34,12 +37,16 @@
 from esbmc_ai.chat_response import FinishReason, ChatResponse
 from esbmc_ai.ai_models import _ai_model_names
 
-
-commands: list[ChatCommand] = []
-command_names: list[str]
 help_command: HelpCommand = HelpCommand()
 fix_code_command: FixCodeCommand = FixCodeCommand()
 exit_command: ExitCommand = ExitCommand()
+command_runner: CommandRunner = CommandRunner(
+    [
+        help_command,
+        exit_command,
+        fix_code_command,
+    ]
+)
 
 chat: UserChat
 
@@ -105,21 +112,11 @@ def print_assistant_response(
         )
 
 
-def init_commands_list() -> None:
-    """Setup Help command and commands list."""
-    # Built in commands
-    global help_command
-    commands.extend(
-        [
-            help_command,
-            exit_command,
-            fix_code_command,
-        ]
-    )
-    help_command.set_commands(commands)
-
-    global command_names
-    command_names = [command.command_name for command in commands]
+def init_addons() -> None:
+    command_runner.addon_commands.clear()
+    command_runner.addon_commands.extend(Config.get_value("addon_modules"))
+    if len(command_runner.addon_commands):
+        printv("Addons:\n\t* " + "\t * ".join(command_runner.addon_commands_names))
 
 
 def update_solution(source_code: str) -> None:
@@ -215,23 +212,7 @@ def _run_command_mode(command: ChatCommand, args: argparse.Namespace) -> None:
     sys.exit(0)
 
 
-def parse_command(user_prompt_string: str) -> tuple[str, list[str]]:
-    """Parses a command and returns it based on the command rules outlined in
-    the wiki: https://github.com/Yiannis128/esbmc-ai/wiki/User-Chat-Mode"""
-    regex_pattern: str = (
-        r'\s+(?=(?:[^\\"]*(?:\\.[^\\"]*)*)$)|(?:(?<!\\)".*?(?<!\\)")|(?:\\.)+|\S+'
-    )
-    segments: list[str] = re.findall(regex_pattern, user_prompt_string)
-    parsed_array: list[str] = [segment for segment in segments if segment != " "]
-    # Remove all empty spaces.
-    command: str = parsed_array[0]
-    command_args: list[str] = parsed_array[1:]
-    return command, command_args
-
-
 def main() -> None:
-    init_commands_list()
-
     parser = argparse.ArgumentParser(
         prog="ESBMC-ChatGPT",
         description=HELP_MESSAGE,
@@ -295,11 +276,9 @@ def main() -> None:
     parser.add_argument(
         "-c",
         "--command",
-        choices=command_names,
-        metavar="",
         help="Runs the program in command mode, it will exit after the command ends with an exit code. Options: {"
-        + ", ".join(command_names)
-        + "}",
+        + ", ".join(command_runner.builtin_commands_names)
+        + "}. To see addon commands avaiilable: Run with '-c help'.",
     )
 
     parser.add_argument(
@@ -326,8 +305,8 @@ def main() -> None:
 
     Config.init(args)
     ESBMCUtil.init(Config.get_value("esbmc.path"))
-
     check_health()
+    init_addons()
 
     printv(f"Source code format: {Config.get_value('source_code_format')}")
     printv(f"ESBMC output type: {Config.get_value('esbmc.output_type')}")
@@ -349,12 +328,19 @@ def main() -> None:
     # If not, then continue to user mode.
     if args.command != None:
         command = args.command
+        command_names: list[str] = command_runner.command_names
         if command in command_names:
             print("Running Command:", command)
             for idx, command_name in enumerate(command_names):
                 if command == command_name:
-                    _run_command_mode(command=commands[idx], args=args)
+                    _run_command_mode(command=command_runner.commands[idx], args=args)
             sys.exit(0)
+        else:
+            print(
+                f"Error: Unknown command: {command}. Choose from: "
+                + ", ".join(command_names)
+            )
+            sys.exit(1)
 
     # ===========================================
     # User Mode (Supports only 1 file)
@@ -440,7 +426,7 @@ def main() -> None:
 
         # Check if it is a command, if not, then pass it to the chat interface.
         if user_message.startswith("/"):
-            command, command_args = parse_command(user_message)
+            command, command_args = CommandRunner.parse_command(user_message)
             command = command[1:]  # Remove the /
             if command == fix_code_command.command_name:
                 # Fix Code command
@@ -458,7 +444,7 @@ def main() -> None:
             else:
                 # Commands without parameters or returns are handled automatically.
                 found: bool = False
-                for cmd in commands:
+                for cmd in command_runner.commands:
                     if cmd.command_name == command:
                         found = True
                         cmd.execute()

esbmc_ai/command_runner.py

@@ -0,0 +1,51 @@
+import re
+from esbmc_ai.commands.chat_command import ChatCommand
+from esbmc_ai.commands.help_command import HelpCommand
+
+
+class CommandRunner:
+    """Command runner manages running and storing commands."""
+
+    def __init__(self, builtin_commands: list[ChatCommand]) -> None:
+        self._builtin_commands: list[ChatCommand] = builtin_commands.copy()
+        self._addon_commands: list[ChatCommand] = []
+
+        # Set the help command commands
+        for cmd in self._builtin_commands:
+            if cmd.command_name == "help":
+                assert isinstance(cmd, HelpCommand)
+                cmd.commands = self.commands
+
+    @property
+    def commands(self) -> list[ChatCommand]:
+        return self._builtin_commands + self._addon_commands
+
+    @property
+    def command_names(self) -> list[str]:
+        return [cmd.command_name for cmd in self.commands]
+
+    @property
+    def builtin_commands_names(self) -> list[str]:
+        return [cmd.command_name for cmd in self._builtin_commands]
+
+    @property
+    def addon_commands_names(self) -> list[str]:
+        return [cmd.command_name for cmd in self._addon_commands]
+
+    @property
+    def addon_commands(self) -> list[ChatCommand]:
+        return self._addon_commands
+
+    @staticmethod
+    def parse_command(user_prompt_string: str) -> tuple[str, list[str]]:
+        """Parses a command and returns it based on the command rules outlined in
+        the wiki: https://github.com/Yiannis128/esbmc-ai/wiki/User-Chat-Mode"""
+        regex_pattern: str = (
+            r'\s+(?=(?:[^\\"]*(?:\\.[^\\"]*)*)$)|(?:(?<!\\)".*?(?<!\\)")|(?:\\.)+|\S+'
+        )
+        segments: list[str] = re.findall(regex_pattern, user_prompt_string)
+        parsed_array: list[str] = [segment for segment in segments if segment != " "]
+        # Remove all empty spaces.
+        command: str = parsed_array[0]
+        command_args: list[str] = parsed_array[1:]
+        return command, command_args

esbmc_ai/commands/help_command.py

@@ -14,9 +14,6 @@ def __init__(self) -> None:
             help_message="Print this help message.",
         )
 
-    def set_commands(self, commands: list[ChatCommand]) -> None:
-        self.commands = commands
-
     @override
     def execute(self, **_: Optional[Any]) -> Optional[Any]:
         print()

esbmc_ai/config.py

@@ -1,12 +1,15 @@
 # Author: Yiannis Charalambous 2023
 
+import importlib
+from importlib.machinery import ModuleSpec
 import os
 import sys
 from platform import system as system_name
 from pathlib import Path
 from dotenv import load_dotenv, find_dotenv
 from langchain.schema import HumanMessage
 import tomllib as toml
+from importlib.util import find_spec
 
 from typing import (
     Any,
@@ -19,7 +22,7 @@
 )
 
 from esbmc_ai.chat_response import list_to_base_messages
-from esbmc_ai.logging import set_verbose
+from esbmc_ai.logging import printv, set_verbose
 from .ai_models import *
 from .api_key_collection import APIKeyCollection
 
@@ -34,29 +37,34 @@
 
 
 class ConfigField(NamedTuple):
-    # The name of the config field and also namespace
     name: str
-    # If a default value is supplied, then it can be omitted from the config
+    """The name of the config field and also namespace"""
     default_value: Any
-    # If true, then the default value will be None, so during
-    # validation, if no value is supplied, then None will be the
-    # the default value, instead of failing due to None being the
-    # default value which under normal circumstances means that the
-    # field is not optional.
+    """If a default value is supplied, then it can be omitted from the config.
+    In order to have a "None" default value, default_value_none must be set."""
     default_value_none: bool = False
-
-    # Lambda function to validate if field has a valid value.
-    # Default is identity function which is return true.
+    """If true, then the default value will be None, so during
+    validation, if no value is supplied, then None will be the
+    the default value, instead of failing due to None being the
+    default value which under normal circumstances means that the
+    field is not optional."""
     validate: Callable[[Any], bool] = lambda _: True
-    # Transform the value once loaded, this allows the value to be saved
-    # as a more complex type than that which is represented in the config
-    # file.
+    """Lambda function to validate if field has a valid value.
+    Default is identity function which is return true."""
     on_load: Callable[[Any], Any] = lambda v: v
-    # If defined, will be called and allows to custom load complex types that
-    # may not match 1-1 in the config. The config file passed as a parameter here
-    # is the original, unflattened version. The value returned should be the value
-    # assigned to this field.
+    """Transform the value once loaded, this allows the value to be saved
+    as a more complex type than that which is represented in the config
+    file.
+    
+    Is ignored if on_read is defined."""
     on_read: Optional[Callable[[dict[str, Any]], Any]] = None
+    """If defined, will be called and allows to custom load complex types that
+    may not match 1-1 in the config. The config file passed as a parameter here
+    is the original, unflattened version. The value returned should be the value
+    assigned to this field.
+    
+    This is a more versatile version of on_load. So if this is used, the on_load
+    will be ignored."""
     error_message: Optional[str] = None
 
 
@@ -89,6 +97,38 @@ def _validate_prompt_template(conv: Dict[str, List[Dict]]) -> bool:
     return True
 
 
+def _validate_addon_modules(mods: list[str]) -> bool:
+    """Validates that all values are string."""
+    for m in mods:
+        if not isinstance(m, str):
+            return False
+        spec: Optional[ModuleSpec] = find_spec(m)
+        if spec is None:
+            return False
+    return True
+
+
+def _init_addon_modules(mods: list[str]) -> list:
+    """Will import addon modules that exist and iterate through the exposed
+    attributes, will then get all available ChatCommands and store them."""
+    from esbmc_ai.commands.chat_command import ChatCommand
+
+    result: list[ChatCommand] = []
+    for module_name in mods:
+        try:
+            m = importlib.import_module(module_name)
+            for attr_name in getattr(m, "__all__"):
+                attr_class = getattr(m, attr_name)
+                if issubclass(attr_class, ChatCommand):
+                    result.append(attr_class())
+                    printv(f"Loading addon: {attr_name}")
+        except ModuleNotFoundError as e:
+            print(f"Addon Loader: Could not import module: {module_name}: {e}")
+            sys.exit(1)
+
+    return result
+
+
 class Config:
     api_keys: APIKeyCollection
     raw_conversation: bool = False
@@ -134,6 +174,14 @@ class Config:
             validate=lambda v: isinstance(v, str) and v in ["full", "single"],
             error_message="source_code_format can only be 'full' or 'single'",
         ),
+        # Store as a list of commands
+        ConfigField(
+            name="addon_modules",
+            default_value=[],
+            validate=_validate_addon_modules,
+            on_load=_init_addon_modules,
+            error_message="addon_modules must be a list of Python modules to load",
+        ),
         ConfigField(
             name="esbmc.path",
             default_value=None,
@@ -351,7 +399,7 @@ def _load_envs(cls) -> None:
         3. esbmc-ai.env file in the current directory, moving upwards in the directory tree.
         4. esbmc-ai.env file in $HOME/.config/ for Linux/macOS and %userprofile% for Windows.
 
-        Note: ESBMC_AI_CFG_PATH undergoes tilde user expansion and also environment
+        Note: ESBMCAI_CONFIG_PATH undergoes tilde user expansion and also environment
         variable expansion.
         """
 
@@ -362,7 +410,7 @@ def get_env_vars() -> None:
                 if value != None:
                     values[k] = value
 
-        keys: list[str] = ["OPENAI_API_KEY", "ESBMC_AI_CFG_PATH"]
+        keys: list[str] = ["OPENAI_API_KEY", "ESBMCAI_CONFIG_PATH"]
         values: dict[str, str] = {}
 
         # Load from system env
@@ -405,7 +453,9 @@ def get_env_vars() -> None:
         )
 
         cls.cfg_path = Path(
-            os.path.expanduser(os.path.expandvars(str(os.getenv("ESBMC_AI_CFG_PATH"))))
+            os.path.expanduser(
+                os.path.expandvars(str(os.getenv("ESBMCAI_CONFIG_PATH")))
+            )
         )
 
     @classmethod