feat(slash_command): control system prompts in /workspace (#736)

Co-authored-by: Oli Morris <[email protected]>

codecompanion-workspace.json

@@ -2,11 +2,14 @@
   "name": "CodeCompanion.nvim",
   "version": "1.0.0",
   "workspace_spec": "1.0",
-  "description": "CodeCompanion.nvim is an AI-powered productivity tool integrated into Neovim, designed to enhance the development workflow by seamlessly interacting with various large language models (LLMs). It offers features like inline code transformations, code creation, refactoring, and supports multiple LLMs such as OpenAI, Anthropic, and Google Gemini, among others. With tools for variable management, agents, and custom workflows, CodeCompanion.nvim streamlines coding tasks and facilitates intelligent code assistance directly within the Neovim editor",
+  "system_prompt": "CodeCompanion.nvim is an AI-powered productivity tool integrated into Neovim, designed to enhance the development workflow by seamlessly interacting with various large language models (LLMs). It offers features like inline code transformations, code creation, refactoring, and supports multiple LLMs such as OpenAI, Anthropic, and Google Gemini, among others. With tools for variable management, agents, and custom workflows, CodeCompanion.nvim streamlines coding tasks and facilitates intelligent code assistance directly within the Neovim editor.",
   "groups": [
     {
       "name": "Chat Buffer",
-      "description": "${workspace_description}. I've grouped a number of files together into a group I'm calling \"${group_name}\". The chat buffer is a Neovim buffer which allows a user to interact with an LLM. The buffer is formatted as Markdown with a user's content residing under a H2 header. The user types their message, saves the buffer and the plugin then uses Tree-sitter to parse the buffer, extracting the contents and sending to an adapter which connects to the user's chosen LLM. The response back from the LLM is streamed into the buffer under another H2 header. The user is then free to respond back to the LLM.\n\nBelow are the relevant files which we will be discussing:\n\n${group_files}",
+      "system_prompt": "I've grouped a number of files together into a group I'm calling \"${group_name}\". The chat buffer is a Neovim buffer which allows a user to interact with an LLM. The buffer is formatted as Markdown with a user's content residing under a H2 header. The user types their message, saves the buffer and the plugin then uses Tree-sitter to parse the buffer, extracting the contents and sending to an adapter which connects to the user's chosen LLM. The response back from the LLM is streamed into the buffer under another H2 header. The user is then free to respond back to the LLM.\n\nBelow are the relevant files which we will be discussing:\n\n${group_files}",
+      "opts": {
+        "remove_config_system_prompt": true
+      },
       "vars": {
         "base_dir": "lua/codecompanion/strategies/chat"
       },
@@ -26,27 +29,6 @@
           "path": "${base_dir}/watchers.lua"
         }
       ]
-    },
-    {
-      "name": "Test",
-      "description": "This is a test group",
-      "vars": {
-        "base_dir": "tests/stubs"
-      },
-      "files": [
-        {
-          "description": "Test description for the file ${filename} located at ${path}",
-          "path": "${base_dir}/stub.go"
-        },
-        "${base_dir}/stub.txt"
-      ],
-      "symbols": [
-        {
-          "description": "Test symbol description for the file ${filename} located at ${path}",
-          "path": "${base_dir}/stub.lua"
-        },
-        "${base_dir}/stub.py"
-      ]
     }
   ]
 }

doc/extending/workspace.md

@@ -4,18 +4,22 @@ Workspaces act as a context management system for your project. This context sit
 
 ## Structure
 
-Taking CodeCompanion's own workspace file as an example:
+Below is an example workspace file for this plugin:
 
 ```json
 {
   "name": "CodeCompanion.nvim",
   "version": "1.0.0",
   "workspace_spec": "1.0",
-  "description": "CodeCompanion.nvim is an AI-powered productivity tool integrated into Neovim, designed to enhance the development workflow by seamlessly interacting with various large language models (LLMs). It offers features like inline code transformations, code creation, refactoring, and supports multiple LLMs such as OpenAI, Anthropic, and Google Gemini, among others. With tools for variable management, agents, and custom workflows, CodeCompanion.nvim streamlines coding tasks and facilitates intelligent code assistance directly within the Neovim editor",
+  "description": "An example workspace file",
+  "system_prompt": "CodeCompanion.nvim is an AI-powered productivity tool integrated into Neovim, designed to enhance the development workflow by seamlessly interacting with various large language models (LLMs). It offers features like inline code transformations, code creation, refactoring, and supports multiple LLMs such as OpenAI, Anthropic, and Google Gemini, among others. With tools for variable management, agents, and custom workflows, CodeCompanion.nvim streamlines coding tasks and facilitates intelligent code assistance directly within the Neovim editor",
   "groups": [
     {
       "name": "Chat Buffer",
-      "description": "...",
+      "system_prompt": "...",
+      "opts": {
+        "remove_config_system_prompt": true
+      },
       "files": [
         {
           "description": "...",
@@ -33,7 +37,10 @@ Taking CodeCompanion's own workspace file as an example:
 }
 ```
 
-- The `description` value contains the high-level description of the project. This is **not** sent to the LLM by default
+- The `description` value contains the high-level description of the workspace file. This is **not** sent to the LLM by default
+- The `system_prompt` value contains text that will be sent to the LLM as a system prompt
+- The `remove_config_system_prompt` key ensures the plugin's default system prompt (as defined in the user's config) is
+removed from the chat buffer
 - The `groups` array contains the grouping of files and symbols that can be shared with the LLM. In this example we just have one group, the _Chat Buffer_
 - The `version` and `workspace_spec` are currently unused
 
@@ -47,7 +54,11 @@ Groups are the core of the workspace file. They are where logical groupings of f
 ```json
 {
   "name": "Chat Buffer",
-  "description": "${workspace_description}. I've grouped a number of files together into a group I'm calling \"${group_name}\". The chat buffer is a Neovim buffer which allows a user to interact with an LLM. The buffer is formatted as Markdown with a user's content residing under a H2 header. The user types their message, saves the buffer and the plugin then uses Tree-sitter to parse the buffer, extracting the contents and sending to an adapter which connects to the user's chosen LLM. The response back from the LLM is streamed into the buffer under another H2 header. The user is then free to respond back to the LLM.\n\nBelow are the relevant files which we will be discussing:\n\n${group_files}",
+  "system_prompt": "I've grouped a number of files together into a group I'm calling \"${group_name}\". The chat buffer is a Neovim buffer which allows a user to interact with an LLM. The buffer is formatted as Markdown with a user's content residing under a H2 header. The user types their message, saves the buffer and the plugin then uses Tree-sitter to parse the buffer, extracting the contents and sending to an adapter which connects to the user's chosen LLM. The response back from the LLM is streamed into the buffer under another H2 header. The user is then free to respond back to the LLM.\n\nBelow are the relevant files which we will be discussing:\n\n${group_files}",
+  "description": "You could also add a description here which will be added as a user prompt",
+  "opts": {
+    "remove_config_system_prompt": true
+  },
   "vars": {
     "base_dir": "lua/codecompanion/strategies/chat"
   },
@@ -72,8 +83,8 @@ Groups are the core of the workspace file. They are where logical groupings of f
 
 There's a lot going on in there:
 
-- Firstly, the `${workspace_description}` is a way of including the description from the top of the workspace file
-- The `${group_name}` provides the name of the current group
+- Firstly, the `system_prompt` within the group is a way of adding to the main, workspace system prompt
+- The `${group_name}` variable provides the name of the current group
 - The `${group_files}` variable contains a list of all the files and symbols in the group
 - The `vars` object is a way of creating variables that can be referenced throughout the files and symbols arrays
 - Each object in the files/symbols array can be a string which defaults to a path, or can be an object containing a
@@ -94,6 +105,7 @@ During conversation with the LLM, it can be useful to also tag the `@files` tool
 A list of all the variables available in workspace files:
 
 - `${workspace_description}` - The description at the top of the workspace file
+- `${workspace_name}` - The name of the workspace file
 - `${group_name}` - The name of the group that is being processed by the slash command
 - `${group_files}` - A list of all the files and symbols in the group
 - `${filename}` - The name of the current file/symbol that is being processed

lua/codecompanion/strategies/chat/init.lua

@@ -420,13 +420,14 @@ function Chat:complete_models(request, callback)
 end
 
 ---Set the system prompt in the chat buffer
+---@prompt? string
 ---@return CodeCompanion.Chat
-function Chat:set_system_prompt()
+function Chat:set_system_prompt(prompt)
   if self.opts and self.opts.ignore_system_prompt then
     return self
   end
 
-  local prompt = config.opts.system_prompt
+  prompt = prompt or config.opts.system_prompt
   if prompt ~= "" then
     if type(prompt) == "function" then
       prompt = prompt({
@@ -459,6 +460,14 @@ function Chat:toggle_system_prompt()
   end
 end
 
+---Remove the system prompt from the chat buffer
+---@return nil
+function Chat:remove_system_prompt()
+  if self.messages[1] and self.messages[1].role == config.constants.SYSTEM_ROLE then
+    table.remove(self.messages, 1)
+  end
+end
+
 ---Parse the last message for any variables
 ---@param message table
 ---@return CodeCompanion.Chat

lua/codecompanion/strategies/chat/slash_commands/workspace.lua

@@ -38,23 +38,29 @@ local function get_file_list(group)
   return table.concat(items, "\n")
 end
 
----Add the description of the group to the chat buffer
----@param chat CodeCompanion.Chat
+---Replace variables in a string
 ---@param workspace table
----@param group { name: string, description: string, files: table?, symbols: table? }
-local function add_group_description(chat, workspace, group)
+---@param group table
+---@param str string
+---@return string
+local function replace_vars(workspace, group, str)
   local builtin = {
     group_name = group.name,
     group_files = get_file_list(group),
     workspace_description = workspace.description,
     workspace_name = workspace.name,
   }
+  return util.replace_placeholders(str, builtin)
+end
 
-  local description = util.replace_placeholders(group.description, builtin)
-
+---Add the description of the group to the chat buffer
+---@param chat CodeCompanion.Chat
+---@param workspace table
+---@param group { name: string, description: string, files: table?, symbols: table? }
+local function add_group_description(chat, workspace, group)
   chat:add_message({
     role = config.constants.USER_ROLE,
-    content = description,
+    content = replace_vars(workspace, group, group.description),
   }, { visible = false })
 end
 
@@ -67,10 +73,11 @@ function SlashCommand.new(args)
     Chat = args.Chat,
     config = args.config,
     context = args.context,
-    workspace = args.workspace or {},
-    opts = args.opts,
+    opts = args.opts or {},
   }, { __index = SlashCommand })
 
+  self.workspace = {}
+
   return self
 end
 
@@ -148,9 +155,27 @@ function SlashCommand:output(selected_group, opts)
     return g.name == selected_group
   end, self.workspace.groups)[1]
 
-  --TODO: Account for all groups
+  if group.opts then
+    if group.opts.remove_config_system_prompt then
+      self.Chat:remove_system_prompt()
+    end
+  end
+
+  -- Add the system prompts
+  if self.workspace.system_prompt or group.system_prompt then
+    local system_prompt = ""
+    if self.workspace.system_prompt and group.system_prompt then
+      system_prompt = self.workspace.system_prompt .. "\n\n" .. group.system_prompt
+    else
+      system_prompt = self.workspace.system_prompt or group.system_prompt
+    end
+    self.Chat:set_system_prompt(replace_vars(self.workspace, group, system_prompt))
+  end
 
-  add_group_description(self.Chat, self.workspace, group)
+  -- Add the description as a user message
+  if group.description then
+    add_group_description(self.Chat, self.workspace, group)
+  end
 
   -- Add files
   if group.files and vim.tbl_count(group.files) > 0 then

tests/strategies/chat/slash_commands/test_workspace.lua

@@ -19,31 +19,38 @@ local expect_starts_with = MiniTest.new_expectation(
 
 local chat
 local wks
-local json
+
+local function set_workspace(path)
+  path = path or "tests/stubs/workspace.json"
+  wks.workspace = wks:read_workspace_file(path)
+end
 
 T["Workspace"] = new_set({
   hooks = {
-    pre_once = function()
+    pre_case = function()
       chat, _ = h.setup_chat_buffer()
       wks = workspace.new({
         Chat = chat,
         context = {},
         opts = {},
       })
     end,
-    post_once = function()
+    post_case = function()
       h.teardown_chat_buffer()
     end,
   },
 })
 
 T["Workspace"]["fetches groups"] = function()
-  wks.workspace = wks:read_workspace_file()
+  set_workspace()
 
-  h.eq("Test", wks.workspace.groups[2].name)
+  h.eq("Test", wks.workspace.groups[1].name)
+  h.eq("Test 2", wks.workspace.groups[2].name)
 end
 
 T["Workspace"]["adds files and symbols"] = function()
+  set_workspace()
+
   h.eq(1, #chat.messages)
   wks:output("Test")
 
@@ -59,4 +66,21 @@ T["Workspace"]["adds files and symbols"] = function()
   expect_starts_with("Here is a symbolic outline of the file `tests/stubs/stub.py`", chat.messages[6].content)
 end
 
+T["Workspace"]["can remove the default system prompt"] = function()
+  set_workspace()
+  wks:output("Test 2")
+
+  h.eq("system", chat.messages[1].role)
+  h.eq("Testing to remove the default system prompt", chat.messages[1].content)
+  h.eq("user", chat.messages[2].role)
+end
+
+T["Workspace"]["can add system prompts"] = function()
+  set_workspace("tests/stubs/workspace_system_prompt.json")
+  wks:output("Test")
+
+  h.eq("system", chat.messages[1].role)
+  h.eq("High level system prompt", chat.messages[1].content)
+end
+
 return T

tests/stubs/workspace.json

@@ -0,0 +1,62 @@
+{
+  "name": "CodeCompanion.nvim",
+  "version": "1.0.0",
+  "workspace_spec": "1.0",
+  "description": "test",
+  "groups": [
+    {
+      "name": "Test",
+      "description": "This is a test group",
+      "vars": {
+        "base_dir": "tests/stubs"
+      },
+      "files": [
+        {
+          "description": "Test description for the file ${filename} located at ${path}",
+          "path": "${base_dir}/stub.go"
+        },
+        "${base_dir}/stub.txt"
+      ],
+      "symbols": [
+        {
+          "description": "Test symbol description for the file ${filename} located at ${path}",
+          "path": "${base_dir}/stub.lua"
+        },
+        "${base_dir}/stub.py"
+      ]
+    },
+    {
+      "name": "Test 2",
+      "system_prompt": "Testing to remove the default system prompt",
+      "opts": {
+        "remove_config_system_prompt": true
+      },
+      "vars": {
+        "base_dir": "tests/stubs"
+      },
+      "files": [
+        {
+          "description": "A test description",
+          "path": "${base_dir}/stub.go"
+        }
+      ]
+    },
+    {
+      "name": "Test 3",
+      "description": "system prompt test",
+      "system_prompt": "This is a system prompt ${workspace_description}",
+      "opts": {
+        "remove_config_system_prompt": true
+      },
+      "vars": {
+        "base_dir": "tests/stubs"
+      },
+      "files": [
+        {
+          "description": "A test description",
+          "path": "${base_dir}/stub.go"
+        }
+      ]
+    }
+  ]
+}

tests/stubs/workspace_system_prompt.json

@@ -0,0 +1,24 @@
+{
+  "name": "Testing System Prompts",
+  "version": "1.0.0",
+  "workspace_spec": "1.0",
+  "system_prompt": "High level system prompt",
+  "groups": [
+    {
+      "name": "Test",
+      "opts": {
+        "remove_config_system_prompt": true
+      },
+      "vars": {
+        "base_dir": "tests/stubs"
+      },
+      "files": [
+        {
+          "description": "A test description",
+          "path": "${base_dir}/stub.go"
+        }
+      ]
+    }
+  ]
+}
+