separated into it's own command system-prompt

Author: aponcedeleonch

src/codegate/db/connection.py

@@ -256,8 +256,7 @@ async def add_workspace(self, workspace_name: str) -> Optional[Workspace]:
         It may raise a ValidationError if the workspace name is invalid.
         or a AlreadyExistsError if the workspace already exists.
         """
-        workspace = Workspace(id=str(uuid.uuid4()), name=workspace_name)
-
+        workspace = Workspace(id=str(uuid.uuid4()), name=workspace_name, system_prompt=None)
         sql = text(
             """
             INSERT INTO workspaces (id, name)

src/codegate/pipeline/cli/cli.py

@@ -8,7 +8,7 @@
     PipelineResult,
     PipelineStep,
 )
-from codegate.pipeline.cli.commands import Version, Workspace
+from codegate.pipeline.cli.commands import SystemPrompt, Version, Workspace
 
 HELP_TEXT = """
 ## CodeGate CLI\n
@@ -32,6 +32,7 @@ async def codegate_cli(command):
     available_commands = {
         "version": Version().exec,
         "workspace": Workspace().exec,
+        "system-prompt": SystemPrompt().exec,
     }
     out_func = available_commands.get(command[0])
     if out_func is None:

src/codegate/pipeline/cli/commands.py

@@ -1,5 +1,5 @@
 from abc import ABC, abstractmethod
-from typing import List
+from typing import Awaitable, Callable, Dict, List, Tuple
 
 from pydantic import ValidationError
 
@@ -8,18 +8,29 @@
 from codegate.workspaces import crud
 
 
+class NoFlagValueError(Exception):
+    pass
+
+class NoSubcommandError(Exception):
+    pass
+
 class CodegateCommand(ABC):
     @abstractmethod
     async def run(self, args: List[str]) -> str:
         pass
 
+    @property
+    @abstractmethod
+    def command_name(self) -> str:
+        pass
+
     @property
     @abstractmethod
     def help(self) -> str:
         pass
 
     async def exec(self, args: List[str]) -> str:
-        if args and args[0] == "-h":
+        if len(args) > 0 and args[0] == "-h":
             return self.help
         return await self.run(args)
 
@@ -28,6 +39,10 @@ class Version(CodegateCommand):
     async def run(self, args: List[str]) -> str:
         return f"CodeGate version: {__version__}"
 
+    @property
+    def command_name(self) -> str:
+        return "version"
+
     @property
     def help(self) -> str:
         return (
@@ -38,18 +53,90 @@ def help(self) -> str:
         )
 
 
-class Workspace(CodegateCommand):
+class CodegateCommandSubcommand(CodegateCommand):
+
+    @property
+    @abstractmethod
+    def subcommands(self) -> Dict[str, Callable[[List[str]], Awaitable[str]]]:
+        pass
+
+    @property
+    @abstractmethod
+    def flags(self) -> List[str]:
+        pass
+
+    def _parse_flags_and_subocomand(self, args: List[str]) -> Tuple[Dict[str, str], List[str], str]:
+        i = 0
+        read_flags = {}
+        # Parse all recognized flags at the start
+        while i < len(args):
+            if args[i] in self.flags:
+                flag_name = args[i]
+                if i + 1 >= len(args):
+                    raise NoFlagValueError(f"Flag {flag_name} needs a value, but none provided.")
+                read_flags[flag_name] = args[i+1]
+                i += 2
+            else:
+                # Once we encounter something that's not a recognized flag,
+                # we assume it's the subcommand
+                break
+
+        if i >= len(args):
+            raise NoSubcommandError("No subcommand found after optional flags.")
+
+        subcommand = args[i]
+        i += 1
+
+        # The rest of the arguments after the subcommand
+        rest = args[i:]
+        return read_flags, rest, subcommand
+
+    async def run(self, args: List[str]) -> str:
+        try:
+            flags, rest, subcommand = self._parse_flags_and_subocomand(args)
+        except NoFlagValueError:
+            return (
+                f"Error reading the command. Flag without value found. "
+                f"Use `codegate {self.command_name} -h` to see available subcommands"
+            )
+        except NoSubcommandError:
+            return (
+                f"Submmand not found "
+                f"Use `codegate {self.command_name} -h` to see available subcommands"
+            )
+
+        command_to_execute = self.subcommands.get(subcommand)
+        if command_to_execute is None:
+            return (
+                f"Submmand not found "
+                f"Use `codegate {self.command_name} -h` to see available subcommands"
+            )
+
+        return await command_to_execute(flags, rest)
+
+
+class Workspace(CodegateCommandSubcommand):
 
     def __init__(self):
         self.workspace_crud = crud.WorkspaceCrud()
-        self.commands = {
+
+    @property
+    def command_name(self) -> str:
+        return "workspace"
+
+    @property
+    def flags(self) -> List[str]:
+        return []
+
+    @property
+    def subcommands(self) -> Dict[str, Callable[[List[str]], Awaitable[str]]]:
+        return {
             "list": self._list_workspaces,
             "add": self._add_workspace,
             "activate": self._activate_workspace,
-            "system-prompt": self._add_system_prompt,
         }
 
-    async def _list_workspaces(self, *args: List[str]) -> str:
+    async def _list_workspaces(self, flags: Dict[str, str], args: List[str]) -> str:
         """
         List all workspaces
         """
@@ -62,7 +149,7 @@ async def _list_workspaces(self, *args: List[str]) -> str:
             respond_str += "\n"
         return respond_str
 
-    async def _add_workspace(self, args: List[str]) -> str:
+    async def _add_workspace(self, flags: Dict[str, str], args: List[str]) -> str:
         """
         Add a workspace
         """
@@ -84,7 +171,7 @@ async def _add_workspace(self, args: List[str]) -> str:
 
         return f"Workspace `{new_workspace_name}` has been added"
 
-    async def _activate_workspace(self, args: List[str]) -> str:
+    async def _activate_workspace(self, flags: Dict[str, str], args: List[str]) -> str:
         """
         Activate a workspace
         """
@@ -105,18 +192,60 @@ async def _activate_workspace(self, args: List[str]) -> str:
             return "An error occurred while activating the workspace"
         return f"Workspace **{workspace_name}** has been activated"
 
-    async def _add_system_prompt(self, args: List[str]):
-        if len(args) < 2:
+    @property
+    def help(self) -> str:
+        return (
+            "### CodeGate Workspace\n"
+            "Manage workspaces.\n\n"
+            "**Usage**: `codegate workspace <command> [args]`\n\n"
+            "Available commands:\n"
+            "- `list`: List all workspaces\n"
+            "  - *args*: None\n"
+            "  - **Usage**: `codegate workspace list`\n"
+            "- `add`: Add a workspace\n"
+            "  - *args*:\n"
+            "    - `workspace_name`\n"
+            "  - **Usage**: `codegate workspace add <workspace_name>`\n"
+            "- `activate`: Activate a workspace\n"
+            "  - *args*:\n"
+            "    - `workspace_name`\n"
+            "  - **Usage**: `codegate workspace activate <workspace_name>`\n"
+        )
+
+
+class SystemPrompt(CodegateCommandSubcommand):
+
+    def __init__(self):
+        self.workspace_crud = crud.WorkspaceCrud()
+
+    @property
+    def command_name(self) -> str:
+        return "system-prompt"
+
+    @property
+    def flags(self) -> List[str]:
+        return ["-w"]
+
+    @property
+    def subcommands(self) -> Dict[str, Callable[[List[str]], Awaitable[str]]]:
+        return {
+            "set": self._set_system_prompt,
+        }
+
+    async def _set_system_prompt(self, flags: Dict[str, str], args: List[str]) -> str:
+        if len(args) == 0:
             return (
                 "Please provide a workspace name and a system prompt. "
-                "Use `codegate workspace system-prompt <workspace_name> <system_prompt>`"
+                "Use `codegate workspace system-prompt -w <workspace_name> <system_prompt>`"
             )
 
-        workspace_name = args[0]
-        sys_prompt_lst = args[1:]
+        workspace_name = flags.get("-w")
+        if not workspace_name:
+            active_workspace = await self.workspace_crud.get_active_workspace()
+            workspace_name = active_workspace.name
 
         updated_worksapce = await self.workspace_crud.update_workspace_system_prompt(
-            workspace_name, sys_prompt_lst
+            workspace_name, args
         )
         if not updated_worksapce:
             return (
@@ -128,37 +257,18 @@ async def _add_system_prompt(self, args: List[str]):
             f"updated to:\n```\n{updated_worksapce.system_prompt}\n```"
         )
 
-    async def run(self, args: List[str]) -> str:
-        if not args:
-            return "Please provide a command. Use `codegate workspace -h` to see available commands"
-        command = args[0]
-        command_to_execute = self.commands.get(command)
-        if command_to_execute is not None:
-            return await command_to_execute(args[1:])
-        else:
-            return "Command not found. Use `codegate workspace -h` to see available commands"
-
     @property
     def help(self) -> str:
         return (
-            "### CodeGate Workspace\n"
-            "Manage workspaces.\n\n"
-            "**Usage**: `codegate workspace <command> [args]`\n\n"
+            "### CodeGate System Prompt\n"
+            "Manage the system prompts of workspaces.\n\n"
+            "**Usage**: `codegate system-prompt -w <workspace_name> <command>`\n\n"
+            "*args*:\n"
+            "- `workspace_name`: Optional workspace name. If not specified will use the "
+            "active workspace\n\n"
             "Available commands:\n"
-            "- `list`: List all workspaces\n"
-            "  - *args*: None\n"
-            "  - **Usage**: `codegate workspace list`\n"
-            "- `add`: Add a workspace\n"
+            "- `set`: Set the system prompt of the workspace\n"
             "  - *args*:\n"
-            "    - `workspace_name`\n"
-            "  - **Usage**: `codegate workspace add <workspace_name>`\n"
-            "- `activate`: Activate a workspace\n"
-            "  - *args*:\n"
-            "    - `workspace_name`\n"
-            "  - **Usage**: `codegate workspace activate <workspace_name>`\n"
-            "- `system-prompt`: Modify the system-prompt of a workspace\n"
-            "  - *args*:\n"
-            "    - `workspace_name`\n"
-            "    - `system_prompt`\n"
-            "  - **Usage**: `codegate workspace system-prompt <workspace_name> <system_prompt>`\n"
+            "    - `system_prompt`: The system prompt to set\n"
+            "  - **Usage**: `codegate system-prompt -w <workspace_name> set <system_prompt>`\n"
         )