Transform codegate-version and codegate-workspace to cli like codegate (#633)

* Transform `codegate-version` and `codegate-workspace` to cli like

This PR transforms the existent shortcut commands for the pipleine
and groups them in a single step. Now both commands can be accessed
using:
- `codegate version`
- `codegate workspace`

* Styling and test changes

Author: aponcedeleonch

src/codegate/api/v1.py

@@ -3,10 +3,10 @@
 from fastapi.routing import APIRoute
 
 from codegate.api import v1_models
-from codegate.pipeline.workspace import commands as wscmd
+from codegate.workspaces.crud import WorkspaceCrud
 
 v1 = APIRouter()
-wscrud = wscmd.WorkspaceCrud()
+wscrud = WorkspaceCrud()
 
 
 def uniq_name(route: APIRoute):
@@ -61,9 +61,12 @@ async def create_workspace(request: v1_models.CreateWorkspaceRequest):
     return v1_models.Workspace(name=request.name)
 
 
-
-@v1.delete("/workspaces/{workspace_name}", tags=["Workspaces"],
-           generate_unique_id_function=uniq_name, status_code=204)
+@v1.delete(
+    "/workspaces/{workspace_name}",
+    tags=["Workspaces"],
+    generate_unique_id_function=uniq_name,
+    status_code=204,
+)
 async def delete_workspace(workspace_name: str):
     """Delete a workspace by name."""
     raise NotImplementedError

src/codegate/api/v1_models.py

@@ -9,36 +9,44 @@ class Workspace(pydantic.BaseModel):
     name: str
     is_active: bool
 
+
 class ActiveWorkspace(Workspace):
     # TODO: use a more specific type for last_updated
     last_updated: Any
 
+
 class ListWorkspacesResponse(pydantic.BaseModel):
     workspaces: list[Workspace]
 
     @classmethod
     def from_db_workspaces(
-        cls, db_workspaces: List[db_models.WorkspaceActive])-> "ListWorkspacesResponse":
-        return cls(workspaces=[
-            Workspace(name=ws.name, is_active=ws.active_workspace_id is not None)
-            for ws in db_workspaces])
+        cls, db_workspaces: List[db_models.WorkspaceActive]
+    ) -> "ListWorkspacesResponse":
+        return cls(
+            workspaces=[
+                Workspace(name=ws.name, is_active=ws.active_workspace_id is not None)
+                for ws in db_workspaces
+            ]
+        )
+
 
 class ListActiveWorkspacesResponse(pydantic.BaseModel):
     workspaces: list[ActiveWorkspace]
 
     @classmethod
     def from_db_workspaces(
-        cls, ws: Optional[db_models.ActiveWorkspace]) -> "ListActiveWorkspacesResponse":
+        cls, ws: Optional[db_models.ActiveWorkspace]
+    ) -> "ListActiveWorkspacesResponse":
         if ws is None:
             return cls(workspaces=[])
-        return cls(workspaces=[
-            ActiveWorkspace(name=ws.name,
-                            is_active=True,
-                            last_updated=ws.last_update)
-                            ])
+        return cls(
+            workspaces=[ActiveWorkspace(name=ws.name, is_active=True, last_updated=ws.last_update)]
+        )
+
 
 class CreateWorkspaceRequest(pydantic.BaseModel):
     name: str
 
+
 class ActivateWorkspaceRequest(pydantic.BaseModel):
     name: str

src/codegate/db/connection.py

@@ -30,6 +30,7 @@
 alert_queue = asyncio.Queue()
 fim_cache = FimCache()
 
+
 class DbCodeGate:
     _instance = None
 
@@ -256,8 +257,7 @@ async def add_workspace(self, workspace_name: str) -> Optional[Workspace]:
             """
         )
         try:
-            added_workspace = await self._execute_update_pydantic_model(
-                workspace, sql)
+            added_workspace = await self._execute_update_pydantic_model(workspace, sql)
         except Exception as e:
             logger.error(f"Failed to add workspace: {workspace_name}.", error=str(e))
             return None

src/codegate/pipeline/cli/cli.py

@@ -0,0 +1,92 @@
+import shlex
+
+from litellm import ChatCompletionRequest
+
+from codegate.pipeline.base import (
+    PipelineContext,
+    PipelineResponse,
+    PipelineResult,
+    PipelineStep,
+)
+from codegate.pipeline.cli.commands import Version, Workspace
+
+HELP_TEXT = """
+## CodeGate CLI\n
+**Usage**: `codegate [-h] <command> [args]`\n
+Check the help of each command by running `codegate <command> -h`\n
+Available commands:
+- `version`: Show the version of CodeGate
+- `workspace`: Perform different operations on workspaces
+"""
+
+NOT_FOUND_TEXT = "Command not found. Use `codegate -h` to see available commands."
+
+
+async def codegate_cli(command):
+    """
+    Process the 'codegate' command.
+    """
+    if len(command) == 0:
+        return HELP_TEXT
+
+    available_commands = {
+        "version": Version().exec,
+        "workspace": Workspace().exec,
+        "-h": lambda _: HELP_TEXT,
+    }
+    out_func = available_commands.get(command[0])
+    if out_func is None:
+        return NOT_FOUND_TEXT
+
+    return await out_func(command[1:])
+
+
+class CodegateCli(PipelineStep):
+    """Pipeline step that handles codegate cli."""
+
+    @property
+    def name(self) -> str:
+        """
+        Returns the name of this pipeline step.
+
+        Returns:
+            str: The identifier 'codegate-cli'
+        """
+        return "codegate-cli"
+
+    async def process(
+        self, request: ChatCompletionRequest, context: PipelineContext
+    ) -> PipelineResult:
+        """
+        Checks if the last user message contains "codegate" and process the command.
+        This short-circuits the pipeline if the message is found.
+
+        Args:
+            request (ChatCompletionRequest): The chat completion request to process
+            context (PipelineContext): The current pipeline context
+
+        Returns:
+            PipelineResult: Contains the response if triggered, otherwise continues
+            pipeline
+        """
+        last_user_message = self.get_last_user_message(request)
+
+        if last_user_message is not None:
+            last_user_message_str, _ = last_user_message
+            splitted_message = last_user_message_str.lower().split(" ")
+            # We expect codegate as the first word in the message
+            if splitted_message[0] == "codegate":
+                context.shortcut_response = True
+                args = shlex.split(last_user_message_str)
+                cmd_out = await codegate_cli(args[1:])
+                return PipelineResult(
+                    response=PipelineResponse(
+                        step_name=self.name,
+                        content=cmd_out,
+                        model=request["model"],
+                    ),
+                    context=context,
+                )
+
+        # Fall through
+        return PipelineResult(request=request, context=context)

src/codegate/pipeline/cli/commands.py

@@ -0,0 +1,125 @@
+from abc import ABC, abstractmethod
+from typing import List
+
+from codegate import __version__
+from codegate.workspaces.crud import WorkspaceCrud
+
+
+class CodegateCommand(ABC):
+    @abstractmethod
+    async def run(self, args: List[str]) -> str:
+        pass
+
+    @property
+    @abstractmethod
+    def help(self) -> str:
+        pass
+
+    async def exec(self, args: List[str]) -> str:
+        if args and args[0] == "-h":
+            return self.help
+        return await self.run(args)
+
+
+class Version(CodegateCommand):
+    async def run(self, args: List[str]) -> str:
+        return f"CodeGate version: {__version__}"
+
+    @property
+    def help(self) -> str:
+        return (
+            "### CodeGate Version\n\n"
+            "Prints the version of CodeGate.\n\n"
+            "**Usage**: `codegate version`\n\n"
+            "*args*: None"
+        )
+
+
+class Workspace(CodegateCommand):
+
+    def __init__(self):
+        self.workspace_crud = WorkspaceCrud()
+        self.commands = {
+            "list": self._list_workspaces,
+            "add": self._add_workspace,
+            "activate": self._activate_workspace,
+        }
+
+    async def _list_workspaces(self, *args: List[str]) -> str:
+        """
+        List all workspaces
+        """
+        workspaces = await self.workspace_crud.get_workspaces()
+        respond_str = ""
+        for workspace in workspaces:
+            respond_str += f"- {workspace.name}"
+            if workspace.active_workspace_id:
+                respond_str += " **(active)**"
+            respond_str += "\n"
+        return respond_str
+
+    async def _add_workspace(self, args: List[str]) -> str:
+        """
+        Add a workspace
+        """
+        if args is None or len(args) == 0:
+            return "Please provide a name. Use `codegate-workspace add your_workspace_name`"
+
+        new_workspace_name = args[0]
+        if not new_workspace_name:
+            return "Please provide a name. Use `codegate-workspace add your_workspace_name`"
+
+        workspace_created = await self.workspace_crud.add_workspace(new_workspace_name)
+        if not workspace_created:
+            return (
+                "Something went wrong. Workspace could not be added.\n"
+                "1. Check if the name is alphanumeric and only contains dashes, and underscores.\n"
+                "2. Check if the workspace already exists."
+            )
+        return f"Workspace **{new_workspace_name}** has been added"
+
+    async def _activate_workspace(self, args: List[str]) -> str:
+        """
+        Activate a workspace
+        """
+        if args is None or len(args) == 0:
+            return "Please provide a name. Use `codegate-workspace activate workspace_name`"
+
+        workspace_name = args[0]
+        if not workspace_name:
+            return "Please provide a name. Use `codegate-workspace activate workspace_name`"
+
+        was_activated = await self.workspace_crud.activate_workspace(workspace_name)
+        if not was_activated:
+            return (
+                f"Workspace **{workspace_name}** does not exist or was already active. "
+                f"Use `codegate-workspace add {workspace_name}` to add it"
+            )
+        return f"Workspace **{workspace_name}** has been activated"
+
+    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\n"
+            "Manage workspaces.\n\n"
+            "**Usage**: `codegate workspace <command> [args]`\n\n"
+            "Available commands:\n\n"
+            "- `list`: List all workspaces\n\n"
+            "  - *args*: None\n\n"
+            "- `add`: Add a workspace\n\n"
+            "  - *args*:\n\n"
+            "    - `workspace_name`\n\n"
+            "- `activate`: Activate a workspace\n\n"
+            "  - *args*:\n\n"
+            "    - `workspace_name`"
+        )

src/codegate/pipeline/factory.py

@@ -2,6 +2,7 @@
 
 from codegate.config import Config
 from codegate.pipeline.base import PipelineStep, SequentialPipelineProcessor
+from codegate.pipeline.cli.cli import CodegateCli
 from codegate.pipeline.codegate_context_retriever.codegate import CodegateContextRetriever
 from codegate.pipeline.extract_snippets.extract_snippets import CodeSnippetExtractor
 from codegate.pipeline.extract_snippets.output import CodeCommentStep
@@ -13,8 +14,6 @@
     SecretUnredactionStep,
 )
 from codegate.pipeline.system_prompt.codegate import SystemPrompt
-from codegate.pipeline.version.version import CodegateVersion
-from codegate.pipeline.workspace.workspace import CodegateWorkspace
 
 
 class PipelineFactory:
@@ -28,8 +27,7 @@ def create_input_pipeline(self) -> SequentialPipelineProcessor:
             # and without obfuscating the secrets, we'd leak the secrets during those
             # later steps
             CodegateSecrets(),
-            CodegateVersion(),
-            CodegateWorkspace(),
+            CodegateCli(),
             CodeSnippetExtractor(),
             CodegateContextRetriever(),
             SystemPrompt(Config.get_config().prompts.default_chat),