From de005b4ce0028ff8785483e0620fb1f196276f2d Mon Sep 17 00:00:00 2001 From: Alejandro Ponce Date: Wed, 22 Jan 2025 12:11:38 +0200 Subject: [PATCH 1/2] Rename workspaces system-prompt to custom-instructions Check the [Discord discussion](https://discord.com/channels/1184987096302239844/1317203257051054120/1331520854101856327) for further context. There was a poll conducted and we decided to carry on the renaming --- ..._0956-90d5471db49a_rename_system_prompt.py | 25 ++++++++ src/codegate/api/v1.py | 26 ++++---- src/codegate/api/v1_models.py | 2 +- src/codegate/db/connection.py | 12 ++-- src/codegate/db/models.py | 12 +++- src/codegate/pipeline/cli/cli.py | 4 +- src/codegate/pipeline/cli/commands.py | 59 ++++++++++--------- .../pipeline/system_prompt/codegate.py | 16 ++--- src/codegate/workspaces/crud.py | 12 ++-- .../system_prompt/test_system_prompt.py | 6 +- 10 files changed, 106 insertions(+), 68 deletions(-) create mode 100644 migrations/versions/2025_01_22_0956-90d5471db49a_rename_system_prompt.py diff --git a/migrations/versions/2025_01_22_0956-90d5471db49a_rename_system_prompt.py b/migrations/versions/2025_01_22_0956-90d5471db49a_rename_system_prompt.py new file mode 100644 index 00000000..ad8171ba --- /dev/null +++ b/migrations/versions/2025_01_22_0956-90d5471db49a_rename_system_prompt.py @@ -0,0 +1,25 @@ +"""rename system_prompt + +Revision ID: 90d5471db49a +Revises: 4dec3e456c9e +Create Date: 2025-01-22 09:56:21.520839+00:00 + +""" + +from typing import Sequence, Union + +from alembic import op + +# revision identifiers, used by Alembic. +revision: str = "90d5471db49a" +down_revision: Union[str, None] = "4dec3e456c9e" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + op.execute("ALTER TABLE workspaces RENAME COLUMN system_prompt TO custom_instructions;") + + +def downgrade() -> None: + op.execute("ALTER TABLE workspaces RENAME COLUMN custom_instructions TO system_prompt;") diff --git a/src/codegate/api/v1.py b/src/codegate/api/v1.py index 60452c83..ec111525 100644 --- a/src/codegate/api/v1.py +++ b/src/codegate/api/v1.py @@ -228,12 +228,12 @@ async def get_workspace_messages(workspace_name: str) -> List[Conversation]: @v1.get( - "/workspaces/{workspace_name}/system-prompt", + "/workspaces/{workspace_name}/custom-instructions", tags=["Workspaces"], generate_unique_id_function=uniq_name, ) -async def get_workspace_system_prompt(workspace_name: str) -> v1_models.SystemPrompt: - """Get the system prompt for a workspace.""" +async def get_workspace_custom_instructions(workspace_name: str) -> v1_models.CustomInstructions: + """Get the custom instructions of a workspace.""" try: ws = await wscrud.get_workspace_by_name(workspace_name) except crud.WorkspaceDoesNotExistError: @@ -241,22 +241,24 @@ async def get_workspace_system_prompt(workspace_name: str) -> v1_models.SystemPr except Exception: raise HTTPException(status_code=500, detail="Internal server error") - if ws.system_prompt is None: - return v1_models.SystemPrompt(prompt="") + if ws.custom_instructions is None: + return v1_models.CustomInstructions(prompt="") - return v1_models.SystemPrompt(prompt=ws.system_prompt) + return v1_models.CustomInstructions(prompt=ws.custom_instructions) @v1.put( - "/workspaces/{workspace_name}/system-prompt", + "/workspaces/{workspace_name}/custom-instructions", tags=["Workspaces"], generate_unique_id_function=uniq_name, status_code=204, ) -async def set_workspace_system_prompt(workspace_name: str, request: v1_models.SystemPrompt): +async def set_workspace_custom_instructions( + workspace_name: str, request: v1_models.CustomInstructions +): try: # This already checks if the workspace exists - await wscrud.update_workspace_system_prompt(workspace_name, [request.prompt]) + await wscrud.update_workspace_custom_instructions(workspace_name, [request.prompt]) except crud.WorkspaceDoesNotExistError: raise HTTPException(status_code=404, detail="Workspace does not exist") except Exception: @@ -266,15 +268,15 @@ async def set_workspace_system_prompt(workspace_name: str, request: v1_models.Sy @v1.delete( - "/workspaces/{workspace_name}/system-prompt", + "/workspaces/{workspace_name}/custom-instructions", tags=["Workspaces"], generate_unique_id_function=uniq_name, status_code=204, ) -async def delete_workspace_system_prompt(workspace_name: str): +async def delete_workspace_custom_instructions(workspace_name: str): try: # This already checks if the workspace exists - await wscrud.update_workspace_system_prompt(workspace_name, []) + await wscrud.update_workspace_custom_instructions(workspace_name, []) except crud.WorkspaceDoesNotExistError: raise HTTPException(status_code=404, detail="Workspace does not exist") except Exception: diff --git a/src/codegate/api/v1_models.py b/src/codegate/api/v1_models.py index bae42a46..efdc35ae 100644 --- a/src/codegate/api/v1_models.py +++ b/src/codegate/api/v1_models.py @@ -10,7 +10,7 @@ class Workspace(pydantic.BaseModel): is_active: bool -class SystemPrompt(pydantic.BaseModel): +class CustomInstructions(pydantic.BaseModel): prompt: str diff --git a/src/codegate/db/connection.py b/src/codegate/db/connection.py index dc75bb83..3ca9bd0a 100644 --- a/src/codegate/db/connection.py +++ b/src/codegate/db/connection.py @@ -271,7 +271,7 @@ async def add_workspace(self, workspace_name: str) -> 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, system_prompt=None) + workspace = Workspace(id=str(uuid.uuid4()), name=workspace_name, custom_instructions=None) sql = text( """ INSERT INTO workspaces (id, name) @@ -294,7 +294,7 @@ async def update_workspace(self, workspace: Workspace) -> Workspace: """ UPDATE workspaces SET name = :name, - system_prompt = :system_prompt + custom_instructions = :custom_instructions WHERE id = :id RETURNING * """ @@ -477,7 +477,7 @@ async def get_archived_workspaces(self) -> List[Workspace]: sql = text( """ SELECT - id, name, system_prompt + id, name, custom_instructions FROM workspaces WHERE deleted_at IS NOT NULL ORDER BY deleted_at DESC @@ -490,7 +490,7 @@ async def get_workspace_by_name(self, name: str) -> Optional[Workspace]: sql = text( """ SELECT - id, name, system_prompt + id, name, custom_instructions FROM workspaces WHERE name = :name AND deleted_at IS NULL """ @@ -505,7 +505,7 @@ async def get_archived_workspace_by_name(self, name: str) -> Optional[Workspace] sql = text( """ SELECT - id, name, system_prompt + id, name, custom_instructions FROM workspaces WHERE name = :name AND deleted_at IS NOT NULL """ @@ -531,7 +531,7 @@ async def get_active_workspace(self) -> Optional[ActiveWorkspace]: sql = text( """ SELECT - w.id, w.name, w.system_prompt, s.id as session_id, s.last_update + w.id, w.name, w.custom_instructions, s.id as session_id, s.last_update FROM sessions s INNER JOIN workspaces w ON w.id = s.active_workspace_id """ diff --git a/src/codegate/db/models.py b/src/codegate/db/models.py index ac54eacc..7d3c0d69 100644 --- a/src/codegate/db/models.py +++ b/src/codegate/db/models.py @@ -47,10 +47,18 @@ class Setting(BaseModel): ] +WorskpaceNameStr = Annotated[ + str, + StringConstraints( + strip_whitespace=True, to_lower=True, pattern=r"^[a-zA-Z0-9_-]+$", strict=True + ), +] + + class Workspace(BaseModel): id: str name: WorskpaceNameStr - system_prompt: Optional[str] + custom_instructions: Optional[str] class Session(BaseModel): @@ -99,6 +107,6 @@ class WorkspaceActive(BaseModel): class ActiveWorkspace(BaseModel): id: str name: str - system_prompt: Optional[str] + custom_instructions: Optional[str] session_id: str last_update: datetime.datetime diff --git a/src/codegate/pipeline/cli/cli.py b/src/codegate/pipeline/cli/cli.py index bfe2bfda..35054339 100644 --- a/src/codegate/pipeline/cli/cli.py +++ b/src/codegate/pipeline/cli/cli.py @@ -8,7 +8,7 @@ PipelineResult, PipelineStep, ) -from codegate.pipeline.cli.commands import SystemPrompt, Version, Workspace +from codegate.pipeline.cli.commands import CustomInstructions, Version, Workspace HELP_TEXT = """ ## CodeGate CLI\n @@ -32,7 +32,7 @@ async def codegate_cli(command): available_commands = { "version": Version().exec, "workspace": Workspace().exec, - "system-prompt": SystemPrompt().exec, + "custom-instructions": CustomInstructions().exec, } out_func = available_commands.get(command[0]) if out_func is None: diff --git a/src/codegate/pipeline/cli/commands.py b/src/codegate/pipeline/cli/commands.py index 24797be8..2008e105 100644 --- a/src/codegate/pipeline/cli/commands.py +++ b/src/codegate/pipeline/cli/commands.py @@ -355,19 +355,19 @@ def help(self) -> str: ) -class SystemPrompt(CodegateCommandSubcommand): +class CustomInstructions(CodegateCommandSubcommand): def __init__(self): self.workspace_crud = crud.WorkspaceCrud() @property def command_name(self) -> str: - return "system-prompt" + return "custom-instructions" @property def flags(self) -> List[str]: """ - Flags for the system-prompt command. + Flags for the custom-instructions command. -w: Workspace name """ return ["-w"] @@ -375,20 +375,20 @@ def flags(self) -> List[str]: @property def subcommands(self) -> Dict[str, Callable[[List[str]], Awaitable[str]]]: return { - "set": self._set_system_prompt, - "show": self._show_system_prompt, - "reset": self._reset_system_prompt, + "set": self._set_custom_instructions, + "show": self._show_custom_instructions, + "reset": self._reset_custom_instructions, } - async def _set_system_prompt(self, flags: Dict[str, str], args: List[str]) -> str: + async def _set_custom_instructions(self, flags: Dict[str, str], args: List[str]) -> str: """ - Set the system prompt of a workspace + Set the custom instructions of a workspace If a workspace name is not provided, the active workspace is used """ if len(args) == 0: return ( - "Please provide a workspace name and a system prompt. " - "Use `codegate workspace system-prompt -w `" + "Please provide a workspace name and custom instructions to use. " + "Use `codegate workspace custom-instructions -w `" ) workspace_name = flags.get("-w") @@ -397,19 +397,20 @@ async def _set_system_prompt(self, flags: Dict[str, str], args: List[str]) -> st workspace_name = active_workspace.name try: - updated_worksapce = await self.workspace_crud.update_workspace_system_prompt( + updated_worksapce = await self.workspace_crud.update_workspace_instructions( workspace_name, args ) except crud.WorkspaceDoesNotExistError: return ( - f"Workspace system prompt not updated. Workspace `{workspace_name}` doesn't exist" + f"Workspace custom instructions not updated. " + f"Workspace `{workspace_name}` doesn't exist" ) - return f"Workspace `{updated_worksapce.name}` system prompt updated." + return f"Workspace `{updated_worksapce.name}` custom instructions updated." - async def _show_system_prompt(self, flags: Dict[str, str], args: List[str]) -> str: + async def _show_custom_instructions(self, flags: Dict[str, str], args: List[str]) -> str: """ - Show the system prompt of a workspace + Show the custom instructions of a workspace If a workspace name is not provided, the active workspace is used """ workspace_name = flags.get("-w") @@ -422,15 +423,15 @@ async def _show_system_prompt(self, flags: Dict[str, str], args: List[str]) -> s except crud.WorkspaceDoesNotExistError: return f"Workspace `{workspace_name}` doesn't exist" - sysprompt = workspace.system_prompt + sysprompt = workspace.custom_instructions if not sysprompt: - return f"Workspace **{workspace.name}** system prompt is unset." + return f"Workspace **{workspace.name}** custom instructions is unset." - return f"Workspace **{workspace.name}** system prompt:\n\n{sysprompt}." + return f"Workspace **{workspace.name}** custom instructions:\n\n{sysprompt}." - async def _reset_system_prompt(self, flags: Dict[str, str], args: List[str]) -> str: + async def _reset_custom_instructions(self, flags: Dict[str, str], args: List[str]) -> str: """ - Reset the system prompt of a workspace + Reset the custom instructions of a workspace If a workspace name is not provided, the active workspace is used """ workspace_name = flags.get("-w") @@ -439,28 +440,28 @@ async def _reset_system_prompt(self, flags: Dict[str, str], args: List[str]) -> workspace_name = active_workspace.name try: - updated_worksapce = await self.workspace_crud.update_workspace_system_prompt( + updated_worksapce = await self.workspace_crud.update_workspace_custom_instructions( workspace_name, [""] ) except crud.WorkspaceDoesNotExistError: return f"Workspace `{workspace_name}` doesn't exist" - return f"Workspace `{updated_worksapce.name}` system prompt reset." + return f"Workspace `{updated_worksapce.name}` custom instructions reset." @property def help(self) -> str: return ( - "### CodeGate System Prompt\n" - "Manage the system prompts of workspaces.\n\n" - "*Note*: If you want to update the system prompt using files please go to the " + "### CodeGate Custom Instructions\n" + "Manage the custom instructionss of workspaces.\n\n" + "*Note*: If you want to update the custom instructions using files please go to the " "[dashboard](http://localhost:9090).\n\n" - "**Usage**: `codegate system-prompt -w `\n\n" + "**Usage**: `codegate custom-instructions -w `\n\n" "*args*:\n" "- `workspace_name`: Optional workspace name. If not specified will use the " "active workspace\n\n" "Available commands:\n" - "- `set`: Set the system prompt of the workspace\n" + "- `set`: Set the custom instructions of the workspace\n" " - *args*:\n" - " - `system_prompt`: The system prompt to set\n" - " - **Usage**: `codegate system-prompt -w set `\n" + " - `instructions`: The custom instructions to set\n" + " - **Usage**: `codegate custom-instructions -w set `\n" ) diff --git a/src/codegate/pipeline/system_prompt/codegate.py b/src/codegate/pipeline/system_prompt/codegate.py index 00efaa0c..76bcf9d1 100644 --- a/src/codegate/pipeline/system_prompt/codegate.py +++ b/src/codegate/pipeline/system_prompt/codegate.py @@ -26,17 +26,17 @@ def name(self) -> str: """ return "system-prompt" - async def _get_workspace_system_prompt(self) -> str: + async def _get_workspace_custom_instructions(self) -> str: wksp_crud = WorkspaceCrud() workspace = await wksp_crud.get_active_workspace() if not workspace: return "" - return workspace.system_prompt + return workspace.custom_instructions async def _construct_system_prompt( self, - wrksp_sys_prompt: str, + wrksp_custom_instr: str, req_sys_prompt: Optional[str], should_add_codegate_sys_prompt: bool, ) -> ChatCompletionSystemMessage: @@ -52,8 +52,8 @@ def _start_or_append(existing_prompt: str, new_prompt: str) -> str: 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) + if wrksp_custom_instr: + system_prompt = _start_or_append(system_prompt, wrksp_custom_instr) # Add request system prompt if present if req_sys_prompt and "codegate" not in req_sys_prompt.lower(): @@ -72,12 +72,12 @@ async def process( to the existing system prompt """ - wrksp_sys_prompt = await self._get_workspace_system_prompt() + wrksp_custom_instructions = await self._get_workspace_custom_instructions() 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: + if not should_add_codegate_sys_prompt and not wrksp_custom_instructions: return PipelineResult(request=request, context=context) new_request = request.copy() @@ -92,7 +92,7 @@ async def process( req_sys_prompt = request_system_message.get("content") system_prompt = await self._construct_system_prompt( - wrksp_sys_prompt, req_sys_prompt, should_add_codegate_sys_prompt + wrksp_custom_instructions, req_sys_prompt, should_add_codegate_sys_prompt ) context.add_alert(self.name, trigger_string=system_prompt) if not request_system_message: diff --git a/src/codegate/workspaces/crud.py b/src/codegate/workspaces/crud.py index 7ad9fe54..7e0daa50 100644 --- a/src/codegate/workspaces/crud.py +++ b/src/codegate/workspaces/crud.py @@ -65,7 +65,9 @@ async def rename_workspace(self, old_workspace_name: str, new_workspace_name: st if not ws: raise WorkspaceDoesNotExistError(f"Workspace {old_workspace_name} does not exist.") db_recorder = DbRecorder() - new_ws = Workspace(id=ws.id, name=new_workspace_name, system_prompt=ws.system_prompt) + new_ws = Workspace( + id=ws.id, name=new_workspace_name, custom_instructions=ws.custom_instructions + ) workspace_renamed = await db_recorder.update_workspace(new_ws) return workspace_renamed @@ -133,18 +135,18 @@ async def recover_workspace(self, workspace_name: str): await db_recorder.recover_workspace(selected_workspace) return - async def update_workspace_system_prompt( - self, workspace_name: str, sys_prompt_lst: List[str] + async def update_workspace_custom_instructions( + self, workspace_name: str, custom_instr_lst: List[str] ) -> Workspace: selected_workspace = await self._db_reader.get_workspace_by_name(workspace_name) if not selected_workspace: raise WorkspaceDoesNotExistError(f"Workspace {workspace_name} does not exist.") - system_prompt = " ".join(sys_prompt_lst) + custom_instructions = " ".join(custom_instr_lst) workspace_update = Workspace( id=selected_workspace.id, name=selected_workspace.name, - system_prompt=system_prompt, + custom_instructions=custom_instructions, ) db_recorder = DbRecorder() updated_workspace = await db_recorder.update_workspace(workspace_update) diff --git a/tests/pipeline/system_prompt/test_system_prompt.py b/tests/pipeline/system_prompt/test_system_prompt.py index f17735e6..6ea36a93 100644 --- a/tests/pipeline/system_prompt/test_system_prompt.py +++ b/tests/pipeline/system_prompt/test_system_prompt.py @@ -29,7 +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="") + step._get_workspace_custom_instructions = AsyncMock(return_value="") # Mock the get_last_user_message method step.get_last_user_message = Mock(return_value=(user_message, 0)) @@ -63,7 +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="") + step._get_workspace_custom_instructions = AsyncMock(return_value="") # Mock the get_last_user_message method step.get_last_user_message = Mock(return_value=(user_message, 0)) @@ -98,7 +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="") + step._get_workspace_custom_instructions = AsyncMock(return_value="") # Mock get_last_user_message to return None step.get_last_user_message = Mock(return_value=None) From 160475a26e120477a476b5a462f29305bde32047 Mon Sep 17 00:00:00 2001 From: Alejandro Ponce Date: Wed, 22 Jan 2025 12:18:44 +0200 Subject: [PATCH 2/2] Fix unit tests --- src/codegate/db/models.py | 8 -------- tests/pipeline/workspace/test_workspace.py | 2 +- 2 files changed, 1 insertion(+), 9 deletions(-) diff --git a/src/codegate/db/models.py b/src/codegate/db/models.py index 7d3c0d69..c4c425da 100644 --- a/src/codegate/db/models.py +++ b/src/codegate/db/models.py @@ -47,14 +47,6 @@ class Setting(BaseModel): ] -WorskpaceNameStr = Annotated[ - str, - StringConstraints( - strip_whitespace=True, to_lower=True, pattern=r"^[a-zA-Z0-9_-]+$", strict=True - ), -] - - class Workspace(BaseModel): id: str name: WorskpaceNameStr diff --git a/tests/pipeline/workspace/test_workspace.py b/tests/pipeline/workspace/test_workspace.py index 891ecf8a..8cadf038 100644 --- a/tests/pipeline/workspace/test_workspace.py +++ b/tests/pipeline/workspace/test_workspace.py @@ -81,7 +81,7 @@ async def test_add_workspaces(args, existing_workspaces, expected_message): with patch("codegate.workspaces.crud.WorkspaceCrud", autospec=True) as mock_recorder_cls: mock_recorder = mock_recorder_cls.return_value workspace_commands.workspace_crud = mock_recorder - created_workspace = WorkspaceModel(id="1", name="myworkspace", system_prompt=None) + created_workspace = WorkspaceModel(id="1", name="myworkspace", custom_instructions=None) mock_recorder.add_workspace = AsyncMock(return_value=created_workspace) # Call the method