Skip to content

Commit

Permalink
Implement reseting system-prompt in CLI (#710)
Browse files Browse the repository at this point in the history
This also cleans up the output of the `system-prompt show` sub-command

Signed-off-by: Juan Antonio Osorio <ozz@stacklok.com>
  • Loading branch information
JAORMX authored Jan 22, 2025
1 parent 66bb8e2 commit 3c6209c
Showing 1 changed file with 25 additions and 1 deletion.
26 changes: 25 additions & 1 deletion src/codegate/pipeline/cli/commands.py
Original file line number Diff line number Diff line change
Expand Up @@ -377,6 +377,7 @@ 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,
}

async def _set_system_prompt(self, flags: Dict[str, str], args: List[str]) -> str:
Expand Down Expand Up @@ -421,7 +422,30 @@ 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"

return f"Workspace **{workspace.name}** system prompt:\n\n{workspace.system_prompt}."
sysprompt = workspace.system_prompt
if not sysprompt:
return f"Workspace **{workspace.name}** system prompt is unset."

return f"Workspace **{workspace.name}** system prompt:\n\n{sysprompt}."

async def _reset_system_prompt(self, flags: Dict[str, str], args: List[str]) -> str:
"""
Reset the system prompt of a workspace
If a workspace name is not provided, the active workspace is used
"""
workspace_name = flags.get("-w")
if not workspace_name:
active_workspace = await self.workspace_crud.get_active_workspace()
workspace_name = active_workspace.name

try:
updated_worksapce = await self.workspace_crud.update_workspace_system_prompt(
workspace_name, [""]
)
except crud.WorkspaceDoesNotExistError:
return f"Workspace `{workspace_name}` doesn't exist"

return f"Workspace `{updated_worksapce.name}` system prompt reset."

@property
def help(self) -> str:
Expand Down

0 comments on commit 3c6209c

Please sign in to comment.