feat: Let the user add their own system prompts

Related: #454

This PR is not ready yet. For the moment it adds the system prompts
to DB and associates it to a workspace.

It's missing to use the system prompt and actually send it to the
LLM

Author: aponcedeleonch

migrations/versions/a692c8b52308_add_workspace_system_prompt.py

@@ -0,0 +1,27 @@
+"""add_workspace_system_prompt
+
+Revision ID: a692c8b52308
+Revises: 5c2f3eee5f90
+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'
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    # Add column to workspaces table
+    op.execute("ALTER TABLE workspaces ADD COLUMN system_prompt TEXT DEFAULT NULL;")
+
+
+def downgrade() -> None:
+    op.execute("ALTER TABLE workspaces DROP COLUMN system_prompt;")

src/codegate/db/connection.py

@@ -272,6 +272,20 @@ async def add_workspace(self, workspace_name: str) -> Optional[Workspace]:
             raise AlreadyExistsError(f"Workspace {workspace_name} already exists.")
         return added_workspace
 
+    async def update_workspace(self, workspace: Workspace) -> Optional[Workspace]:
+        sql = text(
+            """
+            UPDATE workspaces SET
+            name = :name,
+            system_prompt = :system_prompt
+            WHERE id = :id
+            RETURNING *
+            """
+        )
+        # We only pass an object to respect the signature of the function
+        updated_workspace = await self._execute_update_pydantic_model(workspace, sql)
+        return updated_workspace
+
     async def update_session(self, session: Session) -> Optional[Session]:
         sql = text(
             """
@@ -382,11 +396,11 @@ async def get_workspaces(self) -> List[WorkspaceActive]:
         workspaces = await self._execute_select_pydantic_model(WorkspaceActive, sql)
         return workspaces
 
-    async def get_workspace_by_name(self, name: str) -> List[Workspace]:
+    async def get_workspace_by_name(self, name: str) -> Optional[Workspace]:
         sql = text(
             """
             SELECT
-                id, name
+                id, name, system_prompt
             FROM workspaces
             WHERE name = :name
             """

src/codegate/db/models.py

@@ -43,12 +43,18 @@ class Setting(BaseModel):
 class Workspace(BaseModel):
     id: str
     name: str
+    system_prompt: Optional[str]
 
     @field_validator("name", mode="plain")
     @classmethod
-    def name_must_be_alphanumeric(cls, value):
+    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
 
 

src/codegate/pipeline/cli/commands.py

@@ -102,6 +102,18 @@ async def _activate_workspace(self, args: List[str]) -> str:
             )
         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)
+        if not updated_worksapce:
+            return (
+                f"Workspace system prompt not updated. "
+                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}```"
+        )
+
     async def run(self, args: List[str]) -> str:
         if not args:
             return "Please provide a command. Use `codegate workspace -h` to see available commands"
@@ -110,6 +122,8 @@ 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

src/codegate/workspaces/crud.py

@@ -53,7 +53,7 @@ async def _is_workspace_active_or_not_exist(
         sessions = await self._db_reader.get_sessions()
         # The current implementation expects only one active session
         if len(sessions) != 1:
-            raise RuntimeError("Something went wrong. No active session found.")
+            raise WorkspaceCrudError("Something went wrong. More than one session found.")
 
         session = sessions[0]
         if session.active_workspace_id == selected_workspace.id:
@@ -77,3 +77,20 @@ async def activate_workspace(self, workspace_name: str) -> bool:
         db_recorder = DbRecorder()
         await db_recorder.update_session(session)
         return True
+
+    async def update_workspace_system_prompt(
+            self, workspace_name: str, sys_prompt_lst: List[str]
+        ) -> Optional[Workspace]:
+        selected_workspace = await self._db_reader.get_workspace_by_name(workspace_name)
+        if not selected_workspace:
+            return None
+
+        system_prompt = " ".join(sys_prompt_lst)
+        workspace_update = Workspace(
+            id=selected_workspace.id,
+            name=selected_workspace.name,
+            system_prompt=system_prompt,
+        )
+        db_recorder = DbRecorder()
+        updated_workspace = await db_recorder.update_workspace(workspace_update)
+        return updated_workspace