Finished functionality to add wrkspace system prompt

Author: aponcedeleonch

migrations/versions/a692c8b52308_add_workspace_system_prompt.py

@@ -5,15 +5,14 @@
 Create Date: 2025-01-17 16:33:58.464223
 
 """
+
 from typing import Sequence, Union
 
 from alembic import op
-import sqlalchemy as sa
-
 
 # revision identifiers, used by Alembic.
-revision: str = 'a692c8b52308'
-down_revision: Union[str, None] = '5c2f3eee5f90'
+revision: str = "a692c8b52308"
+down_revision: Union[str, None] = "5c2f3eee5f90"
 branch_labels: Union[str, Sequence[str], None] = None
 depends_on: Union[str, Sequence[str], None] = None
 

src/codegate/api/v1.py

@@ -59,9 +59,12 @@ async def create_workspace(request: v1_models.CreateWorkspaceRequest):
     except AlreadyExistsError:
         raise HTTPException(status_code=409, detail="Workspace already exists")
     except ValidationError:
-        raise HTTPException(status_code=400,
-                            detail=("Invalid workspace name. "
-                                    "Please use only alphanumeric characters and dashes"))
+        raise HTTPException(
+            status_code=400,
+            detail=(
+                "Invalid workspace name. " "Please use only alphanumeric characters and dashes"
+            ),
+        )
     except Exception:
         raise HTTPException(status_code=500, detail="Internal server error")
 

src/codegate/dashboard/dashboard.py

@@ -20,23 +20,23 @@
 dashboard_router = APIRouter(tags=["Dashboard"])
 db_reader = None
 
+
 def get_db_reader():
     global db_reader
     if db_reader is None:
         db_reader = DbReader()
     return db_reader
 
+
 def fetch_latest_version() -> str:
     url = "https://api.github.com/repos/stacklok/codegate/releases/latest"
-    headers = {
-        "Accept": "application/vnd.github+json",
-        "X-GitHub-Api-Version": "2022-11-28"
-    }
+    headers = {"Accept": "application/vnd.github+json", "X-GitHub-Api-Version": "2022-11-28"}
     response = requests.get(url, headers=headers, timeout=5)
     response.raise_for_status()
     data = response.json()
     return data.get("tag_name", "unknown")
 
+
 @dashboard_router.get("/dashboard/messages")
 def get_messages(db_reader: DbReader = Depends(get_db_reader)) -> List[Conversation]:
     """
@@ -72,14 +72,15 @@ async def stream_sse():
     """
     return StreamingResponse(generate_sse_events(), media_type="text/event-stream")
 
+
 @dashboard_router.get("/dashboard/version")
 def version_check():
     try:
         latest_version = fetch_latest_version()
 
         # normalize the versions as github will return them with a 'v' prefix
-        current_version = __version__.lstrip('v')
-        latest_version_stripped = latest_version.lstrip('v')
+        current_version = __version__.lstrip("v")
+        latest_version_stripped = latest_version.lstrip("v")
 
         is_latest: bool = latest_version_stripped == current_version
 
@@ -95,15 +96,15 @@ def version_check():
             "current_version": __version__,
             "latest_version": "unknown",
             "is_latest": None,
-            "error": "An error occurred while fetching the latest version"
+            "error": "An error occurred while fetching the latest version",
         }
     except Exception as e:
         logger.error(f"Unexpected error: {str(e)}")
         return {
             "current_version": __version__,
             "latest_version": "unknown",
             "is_latest": None,
-            "error": "An unexpected error occurred"
+            "error": "An unexpected error occurred",
         }
 
 

src/codegate/db/connection.py

@@ -30,9 +30,11 @@
 alert_queue = asyncio.Queue()
 fim_cache = FimCache()
 
+
 class AlreadyExistsError(Exception):
     pass
 
+
 class DbCodeGate:
     _instance = None
 
@@ -266,7 +268,8 @@ async def add_workspace(self, workspace_name: str) -> Optional[Workspace]:
 
         try:
             added_workspace = await self._execute_update_pydantic_model(
-                workspace, sql, should_raise=True)
+                workspace, sql, should_raise=True
+            )
         except IntegrityError as e:
             logger.debug(f"Exception type: {type(e)}")
             raise AlreadyExistsError(f"Workspace {workspace_name} already exists.")
@@ -424,7 +427,7 @@ async def get_active_workspace(self) -> Optional[ActiveWorkspace]:
         sql = text(
             """
             SELECT
-                w.id, w.name, s.id as session_id, s.last_update
+                w.id, w.name, w.system_prompt, s.id as session_id, s.last_update
             FROM sessions s
             INNER JOIN workspaces w ON w.id = s.active_workspace_id
             """

src/codegate/db/models.py

@@ -50,11 +50,6 @@ class Workspace(BaseModel):
     def validate_name(cls, value):
         if not re.match(r"^[a-zA-Z0-9_-]+$", value):
             raise ValueError("name must be alphanumeric and can only contain _ and -")
-        # Avoid workspace names that are the same as commands that way we can do stuff like
-        # `codegate workspace list` and
-        # `codegate workspace my-ws system-prompt` without any conflicts
-        elif value in ["list", "add", "activate", "system-prompt"]:
-            raise ValueError("name cannot be the same as a command")
         return value
 
 
@@ -104,5 +99,6 @@ class WorkspaceActive(BaseModel):
 class ActiveWorkspace(BaseModel):
     id: str
     name: str
+    system_prompt: Optional[str]
     session_id: str
     last_update: datetime.datetime

src/codegate/pipeline/cli/commands.py

@@ -31,10 +31,10 @@ async def run(self, args: List[str]) -> str:
     @property
     def help(self) -> str:
         return (
-            "### CodeGate Version\n\n"
+            "### CodeGate Version\n"
             "Prints the version of CodeGate.\n\n"
+            "*args*: None\n\n"
             "**Usage**: `codegate version`\n\n"
-            "*args*: None"
         )
 
 
@@ -46,6 +46,7 @@ def __init__(self):
             "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:
@@ -66,52 +67,63 @@ 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`"
+            return "Please provide a name. Use `codegate workspace add <workspace_name>`"
 
         new_workspace_name = args[0]
         if not new_workspace_name:
-            return "Please provide a name. Use `codegate workspace add your_workspace_name`"
+            return "Please provide a name. Use `codegate workspace add <workspace_name>`"
 
         try:
             _ = await self.workspace_crud.add_workspace(new_workspace_name)
         except ValidationError:
             return "Invalid workspace name: It should be alphanumeric and dashes"
         except AlreadyExistsError:
-            return f"Workspace **{new_workspace_name}** already exists"
+            return f"Workspace `{new_workspace_name}` already exists"
         except Exception:
             return "An error occurred while adding the workspace"
 
-        return f"Workspace **{new_workspace_name}** has been added"
+        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`"
+            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`"
+            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"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"
+        return f"Workspace `{workspace_name}` has been activated"
 
-    async def _add_system_prompt(self, workspace_name: str, sys_prompt_lst: List[str]):
-        updated_worksapce = await self.workspace_crud.update_workspace_system_prompt(workspace_name, sys_prompt_lst)
+    async def _add_system_prompt(self, args: List[str]):
+        if len(args) < 2:
+            return (
+                "Please provide a workspace name and a system prompt. "
+                "Use `codegate workspace system-prompt <workspace_name> <system_prompt>`"
+            )
+
+        workspace_name = args[0]
+        sys_prompt_lst = args[1:]
+
+        updated_worksapce = await self.workspace_crud.update_workspace_system_prompt(
+            workspace_name, sys_prompt_lst
+        )
         if not updated_worksapce:
             return (
                 f"Workspace system prompt not updated. "
-                f"Check if the workspace **{workspace_name}** exists"
+                f"Check if the workspace `{workspace_name}` exists"
             )
         return (
-            f"Workspace **{updated_worksapce.name}** system prompt "
-            f"updated to:\n\n```{updated_worksapce.system_prompt}```"
+            f"Workspace `{updated_worksapce.name}` system prompt "
+            f"updated to:\n```\n{updated_worksapce.system_prompt}\n```"
         )
 
     async def run(self, args: List[str]) -> str:
@@ -122,23 +134,29 @@ async def run(self, args: List[str]) -> str:
         if command_to_execute is not None:
             return await command_to_execute(args[1:])
         else:
-            if len(args) >= 2 and args[1] == "system-prompt":
-                return await self._add_system_prompt(args[0], args[2:])
             return "Command not found. Use `codegate workspace -h` to see available commands"
 
     @property
     def help(self) -> str:
         return (
-            "### CodeGate Workspace\n\n"
+            "### CodeGate Workspace\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`"
+            "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"
+            "- `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"
         )