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/db/connection.py

@@ -436,7 +436,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,33 +67,33 @@ 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>`"
 
         try:
             await self.workspace_crud.activate_workspace(workspace_name)
@@ -104,16 +105,27 @@ 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, 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:
@@ -124,23 +136,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"
         )

src/codegate/pipeline/system_prompt/codegate.py

@@ -1,4 +1,4 @@
-import json
+from typing import Optional
 
 from litellm import ChatCompletionRequest, ChatCompletionSystemMessage
 
@@ -7,6 +7,7 @@
     PipelineResult,
     PipelineStep,
 )
+from codegate.workspaces.crud import WorkspaceCrud
 
 
 class SystemPrompt(PipelineStep):
@@ -16,7 +17,7 @@ class SystemPrompt(PipelineStep):
     """
 
     def __init__(self, system_prompt: str):
-        self._system_message = ChatCompletionSystemMessage(content=system_prompt, role="system")
+        self.codegate_system_prompt = system_prompt
 
     @property
     def name(self) -> str:
@@ -25,6 +26,44 @@ def name(self) -> str:
         """
         return "system-prompt"
 
+    async def _get_workspace_system_prompt(self) -> str:
+        wksp_crud = WorkspaceCrud()
+        workspace = await wksp_crud.get_active_workspace()
+        if not workspace:
+            return ""
+
+        return workspace.system_prompt
+
+    async def _construct_system_prompt(
+        self,
+        wrksp_sys_prompt: str,
+        req_sys_prompt: Optional[str],
+        should_add_codegate_sys_prompt: bool,
+    ) -> ChatCompletionSystemMessage:
+
+        def _start_or_append(existing_prompt: str, new_prompt: str) -> str:
+            if existing_prompt:
+                return existing_prompt + "\n\nHere are additional instructions:\n\n" + new_prompt
+            return new_prompt
+
+        system_prompt = ""
+        # Add codegate system prompt if secrets or bad packages are found at the beginning
+        if should_add_codegate_sys_prompt:
+            system_prompt = _start_or_append(system_prompt, self.codegate_system_prompt)
+
+        # Add workspace system prompt if present
+        if wrksp_sys_prompt:
+            system_prompt = _start_or_append(system_prompt, wrksp_sys_prompt)
+
+        # Add request system prompt if present
+        if req_sys_prompt and "codegate" not in req_sys_prompt.lower():
+            system_prompt = _start_or_append(system_prompt, req_sys_prompt)
+
+        return system_prompt
+
+    async def _should_add_codegate_system_prompt(self, context: PipelineContext) -> bool:
+        return context.secrets_found or context.bad_packages_found
+
     async def process(
         self, request: ChatCompletionRequest, context: PipelineContext
     ) -> PipelineResult:
@@ -33,32 +72,35 @@ async def process(
         to the existing system prompt
         """
 
-        # Nothing to do if no secrets or bad_packages are found
-        if not (context.secrets_found or context.bad_packages_found):
+        wrksp_sys_prompt = await self._get_workspace_system_prompt()
+        should_add_codegate_sys_prompt = await self._should_add_codegate_system_prompt(context)
+
+        # Nothing to do if no secrets or bad_packages are found and we don't have a workspace
+        # system prompt
+        if not should_add_codegate_sys_prompt and not wrksp_sys_prompt:
             return PipelineResult(request=request, context=context)
 
         new_request = request.copy()
 
         if "messages" not in new_request:
             new_request["messages"] = []
 
-        request_system_message = None
+        request_system_message = {}
         for message in new_request["messages"]:
             if message["role"] == "system":
                 request_system_message = message
+        req_sys_prompt = request_system_message.get("content")
 
-        if request_system_message is None:
-            # Add system message
-            context.add_alert(self.name, trigger_string=json.dumps(self._system_message))
-            new_request["messages"].insert(0, self._system_message)
-        elif "codegate" not in request_system_message["content"].lower():
-            # Prepend to the system message
-            prepended_message = (
-                self._system_message["content"]
-                + "\n Here are additional instructions. \n "
-                + request_system_message["content"]
-            )
-            context.add_alert(self.name, trigger_string=prepended_message)
-            request_system_message["content"] = prepended_message
+        system_prompt = await self._construct_system_prompt(
+            wrksp_sys_prompt, req_sys_prompt, should_add_codegate_sys_prompt
+        )
+        context.add_alert(self.name, trigger_string=system_prompt)
+        if not request_system_message:
+            # Insert the system prompt at the beginning of the messages
+            sytem_message = ChatCompletionSystemMessage(content=system_prompt, role="system")
+            new_request["messages"].insert(0, sytem_message)
+        else:
+            # Update the existing system prompt
+            request_system_message["content"] = system_prompt
 
         return PipelineResult(request=new_request, context=context)

src/codegate/workspaces/crud.py

@@ -84,8 +84,8 @@ async def activate_workspace(self, workspace_name: str):
         return
 
     async def update_workspace_system_prompt(
-            self, workspace_name: str, sys_prompt_lst: List[str]
-        ) -> Optional[Workspace]:
+        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

tests/pipeline/system_prompt/test_system_prompt.py

@@ -1,4 +1,4 @@
-from unittest.mock import Mock
+from unittest.mock import AsyncMock, Mock
 
 import pytest
 from litellm.types.llms.openai import ChatCompletionRequest
@@ -14,7 +14,7 @@ def test_init_with_system_message(self):
         """
         test_message = "Test system prompt"
         step = SystemPrompt(system_prompt=test_message)
-        assert step._system_message["content"] == test_message
+        assert step.codegate_system_prompt == test_message
 
     @pytest.mark.asyncio
     async def test_process_system_prompt_insertion(self):
@@ -29,6 +29,7 @@ async def test_process_system_prompt_insertion(self):
         # Create system prompt step
         system_prompt = "Security analysis system prompt"
         step = SystemPrompt(system_prompt=system_prompt)
+        step._get_workspace_system_prompt = AsyncMock(return_value="")
 
         # Mock the get_last_user_message method
         step.get_last_user_message = Mock(return_value=(user_message, 0))
@@ -62,6 +63,7 @@ async def test_process_system_prompt_update(self):
         # Create system prompt step
         system_prompt = "Security analysis system prompt"
         step = SystemPrompt(system_prompt=system_prompt)
+        step._get_workspace_system_prompt = AsyncMock(return_value="")
 
         # Mock the get_last_user_message method
         step.get_last_user_message = Mock(return_value=(user_message, 0))
@@ -74,7 +76,7 @@ async def test_process_system_prompt_update(self):
         assert result.request["messages"][0]["role"] == "system"
         assert (
             result.request["messages"][0]["content"]
-            == system_prompt + "\n Here are additional instructions. \n " + request_system_message
+            == system_prompt + "\n\nHere are additional instructions:\n\n" + request_system_message
         )
         assert result.request["messages"][1]["role"] == "user"
         assert result.request["messages"][1]["content"] == user_message
@@ -96,6 +98,7 @@ async def test_edge_cases(self, edge_case):
 
         system_prompt = "Security edge case prompt"
         step = SystemPrompt(system_prompt=system_prompt)
+        step._get_workspace_system_prompt = AsyncMock(return_value="")
 
         # Mock get_last_user_message to return None
         step.get_last_user_message = Mock(return_value=None)

tests/pipeline/workspace/test_workspace.py

@@ -54,11 +54,11 @@ async def test_list_workspaces(mock_workspaces, expected_output):
     "args, existing_workspaces, expected_message",
     [
         # Case 1: No workspace name provided
-        ([], [], "Please provide a name. Use `codegate workspace add your_workspace_name`"),
+        ([], [], "Please provide a name. Use `codegate workspace add <workspace_name>`"),
         # Case 2: Workspace name is empty string
-        ([""], [], "Please provide a name. Use `codegate workspace add your_workspace_name`"),
+        ([""], [], "Please provide a name. Use `codegate workspace add <workspace_name>`"),
         # Case 3: Successful add
-        (["myworkspace"], [], "Workspace **myworkspace** has been added"),
+        (["myworkspace"], [], "Workspace `myworkspace` has been added"),
     ],
 )
 async def test_add_workspaces(args, existing_workspaces, expected_message):