Add support for #workspace (workspace file map) context

Support workspace map context. This context will simply include all
related filenames to the prompt in chat context. Useful for searching
where something might be (but limited to only searching on filenames
for now).

Signed-off-by: Tomas Slusny <[email protected]>

Author: deathbeam

README.md

@@ -154,6 +154,7 @@ Supported contexts are:
 
 - `buffers` - Includes all open buffers in chat context
 - `buffer` - Includes only the current buffer in chat context
+- `files` - Includes all non-hidden filenames in the current workspace in chat context
 
 ### API
 
@@ -237,7 +238,7 @@ Also see [here](/lua/CopilotChat/config.lua):
   system_prompt = prompts.COPILOT_INSTRUCTIONS, -- System prompt to use
   model = 'gpt-4o', -- Default model to use, see ':CopilotChatModels' for available models
   agent = 'copilot', -- Default agent to use, see ':CopilotChatAgents' for available agents (can be specified manually in prompt via @).
-  context = nil, -- Default context to use, 'buffers', 'buffer' or none (can be specified manually in prompt via #).
+  context = nil, -- Default context to use, 'buffers', 'buffer', 'files' or none (can be specified manually in prompt via #).
   temperature = 0.1, -- GPT result temperature
 
   question_header = '## User ', -- Header to use for user questions

lua/CopilotChat/config.lua

@@ -97,7 +97,7 @@ return {
   system_prompt = prompts.COPILOT_INSTRUCTIONS, -- System prompt to use
   model = 'gpt-4o', -- Default model to use, see ':CopilotChatModels' for available models
   agent = 'copilot', -- Default agent to use, see ':CopilotChatAgents' for available agents (can be specified manually in prompt via @).
-  context = nil, -- Default context to use, 'buffers', 'buffer' or none (can be specified manually in prompt via #).
+  context = nil, -- Default context to use, 'buffers', 'buffer', 'files' or none (can be specified manually in prompt via #).
   temperature = 0.1, -- GPT result temperature
 
   question_header = '## User ', -- Header to use for user questions

lua/CopilotChat/context.lua

@@ -1,3 +1,4 @@
+local async = require('plenary.async')
 local log = require('plenary.log')
 
 local M = {}
@@ -72,6 +73,78 @@ local function data_ranked_by_relatedness(query, data, top_n)
   return result
 end
 
+local get_context_data = async.wrap(function(context, bufnr, callback)
+  vim.schedule(function()
+    local outline = {}
+    if context == 'buffers' then
+      outline = vim.tbl_map(
+        M.build_outline,
+        vim.tbl_filter(function(b)
+          return vim.api.nvim_buf_is_loaded(b) and vim.fn.buflisted(b) == 1
+        end, vim.api.nvim_list_bufs())
+      )
+    elseif context == 'buffer' then
+      table.insert(outline, M.build_outline(bufnr))
+    elseif context == 'files' then
+      outline = M.build_file_map()
+    end
+
+    callback(outline)
+  end)
+end, 3)
+
+--- Get supported contexts
+---@return table<string>
+function M.supported_contexts()
+  return {
+    buffers = 'Includes all open buffers in chat context',
+    buffer = 'Includes only the current buffer in chat context',
+    files = 'Includes all non-hidden filenames in the current workspace in chat context',
+  }
+end
+
+--- Get list of all files in workspace
+---@return table<CopilotChat.copilot.embed>
+function M.build_file_map()
+  -- Use vim.fn.glob() to get all files
+  local files = vim.fn.glob('**/*', false, true)
+
+  -- Filter out hidden files
+  local filtered_files = vim.tbl_filter(function(file)
+    local file_name = vim.fn.fnamemodify(file, ':t')
+    local file_dir = vim.fn.fnamemodify(file, ':h')
+
+    if file_name:match('^%.') or file_dir:match('^%.') then
+      return false
+    end
+
+    return vim.fn.isdirectory(file) == 0
+  end, files)
+
+  if #filtered_files == 0 then
+    return {}
+  end
+
+  local out = {}
+
+  -- Create embeddings in chunks
+  local chunk_size = 100
+  for i = 1, #filtered_files, chunk_size do
+    local chunk = {}
+    for j = i, math.min(i + chunk_size - 1, #filtered_files) do
+      table.insert(chunk, filtered_files[j])
+    end
+
+    table.insert(out, {
+      content = table.concat(chunk, '\n'),
+      filename = 'project_map',
+      filetype = 'text',
+    })
+  end
+
+  return out
+end
+
 --- Build an outline for a buffer
 --- FIXME: Handle multiline function argument definitions when building the outline
 ---@param bufnr number
@@ -198,21 +271,7 @@ function M.find_for_query(copilot, opts)
   local filetype = opts.filetype
   local bufnr = opts.bufnr
 
-  local outline = {}
-  if context == 'buffers' then
-    -- For multiple buffers, only make outlines
-    outline = vim.tbl_map(
-      function(b)
-        return M.build_outline(b)
-      end,
-      vim.tbl_filter(function(b)
-        return vim.api.nvim_buf_is_loaded(b) and vim.fn.buflisted(b) == 1
-      end, vim.api.nvim_list_bufs())
-    )
-  elseif context == 'buffer' then
-    table.insert(outline, M.build_outline(bufnr))
-  end
-
+  local outline = get_context_data(context, bufnr)
   outline = vim.tbl_filter(function(item)
     return item ~= nil
   end, outline)

lua/CopilotChat/copilot.lua

@@ -323,8 +323,8 @@ local function generate_embedding_request(inputs, model)
       if input.content then
         out = out
           .. string.format(
-            'File: `%s`\n```%s\n%s\n```',
-            input.filename,
+            '# FILE:%s CONTEXT\n```%s\n%s\n```',
+            input.filename:upper(),
             input.filetype,
             input.content
           )

lua/CopilotChat/init.lua

@@ -247,6 +247,7 @@ end
 function M.complete_items(callback)
   async.run(function()
     local agents = state.copilot:list_agents()
+    local contexts = context.supported_contexts()
     local items = {}
     local prompts_to_use = M.prompts()
 
@@ -273,23 +274,16 @@ function M.complete_items(callback)
       }
     end
 
-    items[#items + 1] = {
-      word = '#buffers',
-      kind = 'context',
-      menu = 'Include all loaded buffers in context',
-      icase = 1,
-      dup = 0,
-      empty = 0,
-    }
-
-    items[#items + 1] = {
-      word = '#buffer',
-      kind = 'context',
-      menu = 'Include the specified buffer in context',
-      icase = 1,
-      dup = 0,
-      empty = 0,
-    }
+    for prompt_context, description in ipairs(contexts) do
+      items[#items + 1] = {
+        word = '#' .. prompt_context,
+        kind = 'context',
+        menu = description,
+        icase = 1,
+        dup = 0,
+        empty = 0,
+      }
+    end
 
     vim.schedule(function()
       callback(items)
@@ -499,12 +493,13 @@ function M.ask(prompt, config, source)
   append('\n\n' .. config.answer_header .. config.separator .. '\n\n', config)
 
   local selected_context = config.context
-  if string.find(prompt, '#buffers') then
-    selected_context = 'buffers'
-  elseif string.find(prompt, '#buffer') then
-    selected_context = 'buffer'
+  local contexts = vim.tbl_keys(context.supported_contexts())
+  for prompt_context in updated_prompt:gmatch('#([%w_-]+)') do
+    if vim.tbl_contains(contexts, prompt_context) then
+      selected_context = prompt_context
+      updated_prompt = string.gsub(updated_prompt, '#' .. prompt_context .. '%s*', '')
+    end
   end
-  updated_prompt = string.gsub(updated_prompt, '#buffers?%s*', '')
 
   async.run(function()
     local agents = state.copilot:list_agents()