feat(nvim): Allow the codecompanion tool to send chunks instead of full documents. (#180)

* Add an option to return only chunks to LLMs

* refactor(nvim): improve chunk handling and option resolution.

* refactor(nvim): Improve options handling and mode-specific messages

* refactor(nvim): Allow table type for `max_num` in tool options

* docs(nvim): Add type annotations to documentation for setup options

* Auto generate docs

* fix(nvim): fix tool option type conversions

* feat(nvim): clarify line number usage in code chunk descriptions

---------

Co-authored-by: Jacob Segal <jacob.e.segal@gmail.com>
Co-authored-by: Davidyz <Davidyz@users.noreply.github.com>

Author: Davidyz

doc/VectorCode.txt

@@ -167,30 +167,34 @@ This function initialises the VectorCode client and sets up some default
 
 >lua
     -- Default configuration
-    require("vectorcode").setup({
-      cli_cmds = {
-        vectorcode = "vectorcode",
-      },
-      async_opts = {
-        debounce = 10,
-        events = { "BufWritePost", "InsertEnter", "BufReadPost" },
+    require("vectorcode").setup(
+      ---@type VectorCode.Opts
+      {
+        cli_cmds = {
+          vectorcode = "vectorcode",
+        },
+        ---@type VectorCode.RegisterOpts
+        async_opts = {
+          debounce = 10,
+          events = { "BufWritePost", "InsertEnter", "BufReadPost" },
+          exclude_this = true,
+          n_query = 1,
+          notify = false,
+          query_cb = require("vectorcode.utils").make_surrounding_lines_cb(-1),
+          run_on_register = false,
+        },
+        async_backend = "default", -- or "lsp"
         exclude_this = true,
         n_query = 1,
-        notify = false,
-        query_cb = require("vectorcode.utils").make_surrounding_lines_cb(-1),
-        run_on_register = false,
-      },
-      async_backend = "default", -- or "lsp"
-      exclude_this = true,
-      n_query = 1,
-      notify = true,
-      timeout_ms = 5000,
-      on_setup = {
-        update = false, -- set to true to enable update when `setup` is called.
-        lsp = false,
+        notify = true,
+        timeout_ms = 5000,
+        on_setup = {
+          update = false, -- set to true to enable update when `setup` is called.
+          lsp = false,
+        }
+        sync_log_env_var = false,
       }
-      sync_log_env_var = false,
-    })
+    )
 <
 
 The following are the available options for the parameter of this function: -

docs/neovim.md

@@ -141,30 +141,34 @@ This function initialises the VectorCode client and sets up some default
 
 ```lua
 -- Default configuration
-require("vectorcode").setup({
-  cli_cmds = {
-    vectorcode = "vectorcode",
-  },
-  async_opts = {
-    debounce = 10,
-    events = { "BufWritePost", "InsertEnter", "BufReadPost" },
+require("vectorcode").setup(
+  ---@type VectorCode.Opts
+  {
+    cli_cmds = {
+      vectorcode = "vectorcode",
+    },
+    ---@type VectorCode.RegisterOpts
+    async_opts = {
+      debounce = 10,
+      events = { "BufWritePost", "InsertEnter", "BufReadPost" },
+      exclude_this = true,
+      n_query = 1,
+      notify = false,
+      query_cb = require("vectorcode.utils").make_surrounding_lines_cb(-1),
+      run_on_register = false,
+    },
+    async_backend = "default", -- or "lsp"
     exclude_this = true,
     n_query = 1,
-    notify = false,
-    query_cb = require("vectorcode.utils").make_surrounding_lines_cb(-1),
-    run_on_register = false,
-  },
-  async_backend = "default", -- or "lsp"
-  exclude_this = true,
-  n_query = 1,
-  notify = true,
-  timeout_ms = 5000,
-  on_setup = {
-    update = false, -- set to true to enable update when `setup` is called.
-    lsp = false,
+    notify = true,
+    timeout_ms = 5000,
+    on_setup = {
+      update = false, -- set to true to enable update when `setup` is called.
+      lsp = false,
+    }
+    sync_log_env_var = false,
   }
-  sync_log_env_var = false,
-})
+)
 ```
 
 The following are the available options for the parameter of this function:

lua/vectorcode/integrations/codecompanion/common.lua

@@ -3,14 +3,17 @@ local vc_config = require("vectorcode.config")
 local notify_opts = vc_config.notify_opts
 local logger = vc_config.logger
 
----@class VectorCode.CodeCompanion.ToolOpts
----@field max_num integer?
----@field default_num integer?
----@field include_stderr boolean?
----@field use_lsp boolean?
----@field auto_submit table<string, boolean>?
----@field ls_on_start boolean?
----@field no_duplicate boolean?
+---@type VectorCode.CodeCompanion.ToolOpts
+local default_options = {
+  max_num = { chunk = -1, document = -1 },
+  default_num = { chunk = 50, document = 10 },
+  include_stderr = false,
+  use_lsp = false,
+  auto_submit = { ls = false, query = false },
+  ls_on_start = false,
+  no_duplicate = true,
+  chunk_mode = false,
+}
 
 return {
   tool_result_source = "VectorCodeToolResult",
@@ -23,6 +26,68 @@ return {
     return table.concat(vim.iter(t):flatten(math.huge):totable(), "\n")
   end,
 
+  ---@param opts VectorCode.CodeCompanion.ToolOpts|{}|nil
+  ---@return VectorCode.CodeCompanion.ToolOpts
+  get_tool_opts = function(opts)
+    if opts == nil or opts.use_lsp == nil then
+      opts = vim.tbl_deep_extend(
+        "force",
+        opts or {},
+        { use_lsp = vc_config.get_user_config().async_backend == "lsp" }
+      )
+    end
+    opts = vim.tbl_deep_extend("force", default_options, opts)
+    if type(opts.default_num) == "table" then
+      if opts.chunk_mode then
+        opts.default_num = opts.default_num.chunk
+      else
+        opts.default_num = opts.default_num.document
+      end
+      assert(
+        type(opts.default_num) == "number",
+        "default_num should be an integer or a table: {chunk: integer, document: integer}"
+      )
+    end
+    if type(opts.max_num) == "table" then
+      if opts.chunk_mode then
+        opts.max_num = opts.max_num.chunk
+      else
+        opts.max_num = opts.max_num.document
+      end
+      assert(
+        type(opts.max_num) == "number",
+        "max_num should be an integer or a table: {chunk: integer, document: integer}"
+      )
+    end
+    return opts
+  end,
+
+  ---@param result VectorCode.Result
+  ---@return string
+  process_result = function(result)
+    local llm_message
+    if result.chunk then
+      -- chunk mode
+      llm_message =
+        string.format("<path>%s</path><chunk>%s</chunk>", result.path, result.chunk)
+      if result.start_line and result.end_line then
+        llm_message = llm_message
+          .. string.format(
+            "<start_line>%d</start_line><end_line>%d</end_line>",
+            result.start_line,
+            result.end_line
+          )
+      end
+    else
+      -- full document mode
+      llm_message = string.format(
+        "<path>%s</path><content>%s</content>",
+        result.path,
+        result.document
+      )
+    end
+    return llm_message
+  end,
   ---@param use_lsp boolean
   ---@return VectorCode.JobRunner
   initialise_runner = function(use_lsp)

lua/vectorcode/integrations/codecompanion/func_calling_tool.lua

@@ -10,22 +10,19 @@ local job_runner = nil
 ---@param opts VectorCode.CodeCompanion.ToolOpts?
 ---@return CodeCompanion.Agent.Tool
 return check_cli_wrap(function(opts)
-  if opts == nil or opts.use_lsp == nil then
-    opts = vim.tbl_deep_extend(
-      "force",
-      opts or {},
-      { use_lsp = vc_config.get_user_config().async_backend == "lsp" }
-    )
+  opts = cc_common.get_tool_opts(opts)
+  assert(
+    type(opts.max_num) == "number" and type(opts.default_num) == "number",
+    string.format("Options are not correctly formatted:%s", vim.inspect(opts))
+  )
+  ---@type "file"|"chunk"
+  local mode
+  if opts.chunk_mode then
+    mode = "chunk"
+  else
+    mode = "file"
   end
-  opts = vim.tbl_deep_extend("force", {
-    max_num = -1,
-    default_num = 10,
-    include_stderr = false,
-    use_lsp = false,
-    auto_submit = { ls = false, query = false },
-    ls_on_start = false,
-    no_duplicate = true,
-  }, opts or {})
+
   logger.info("Creating CodeCompanion tool with the following args:\n", opts)
   local capping_message = ""
   if opts.max_num > 0 then
@@ -58,7 +55,6 @@ return check_cli_wrap(function(opts)
         end
 
         if action.command == "query" then
-          local args = { "query", "--pipe", "-n", tostring(action.options.count) }
           if action.options.query == nil then
             return {
               status = "error",
@@ -68,7 +64,14 @@ return check_cli_wrap(function(opts)
           if type(action.options.query) == "string" then
             action.options.query = { action.options.query }
           end
+          local args = { "query" }
           vim.list_extend(args, action.options.query)
+          vim.list_extend(args, { "--pipe", "-n", tostring(action.options.count) })
+          if opts.chunk_mode then
+            vim.list_extend(args, { "--include", "path", "chunk" })
+          else
+            vim.list_extend(args, { "--include", "path", "document" })
+          end
           if action.options.project_root == "" then
             action.options.project_root = nil
           end
@@ -188,6 +191,7 @@ return check_cli_wrap(function(opts)
     system_prompt = function()
       local guidelines = {
         "  - The path of a retrieved file will be wrapped in `<path>` and `</path>` tags. Its content will be right after the `</path>` tag, wrapped by `<content>` and `</content>` tags. Do not include the `<path>``</path>` tags in your answers when you mention the paths.",
+        "  - The results may also be chunks of the source code. In this case, the text chunks will be wrapped in <chunk></chunk>. If the starting and ending line ranges are available, they will be wrapped in <start_line></start_line> and <end_line></end_line> tags. Make use of the line numbers (NOT THE XML TAGS) when you're quoting the source code.",
         "  - If you used the tool, tell users that they may need to wait for the results and there will be a virtual text indicator showing the tool is still running",
         "  - Include one single command call for VectorCode each time. You may include multiple keywords in the command",
         "  - VectorCode is the name of this tool. Do not include it in the query unless the user explicitly asks",
@@ -271,43 +275,38 @@ return check_cli_wrap(function(opts)
         if cmd.command == "query" then
           local max_result = #stdout
           if opts.max_num > 0 then
-            max_result = math.min(opts.max_num, max_result)
+            max_result = math.min(opts.max_num or 1, max_result)
           end
           for i, file in pairs(stdout) do
             if i <= max_result then
               if i == 1 then
+                user_message = string.format(
+                  "**VectorCode Tool**: Retrieved %d %s(s)",
+                  max_result,
+                  mode
+                )
                 if cmd.options.project_root then
-                  user_message = string.format(
-                    "**VectorCode Tool**: Retrieved %s files from %s",
-                    max_result,
-                    cmd.options.project_root
-                  )
-                else
-                  user_message =
-                    string.format("**VectorCode Tool**: Retrieved %s files", max_result)
+                  user_message = user_message .. " from " .. cmd.options.project_root
                 end
+                user_message = user_message .. "\n"
               else
                 user_message = ""
               end
-              local llm_message = string.format(
-                [[Here is a file the VectorCode tool retrieved:
-<path>
-%s
-</path>
-<content>
-%s
-</content>
-]],
-                file.path,
-                file.document
+              agent.chat:add_tool_output(
+                self,
+                cc_common.process_result(file),
+                user_message
               )
-              agent.chat:add_tool_output(self, llm_message, user_message)
-              agent.chat.references:add({
-                source = cc_common.tool_result_source,
-                id = file.path,
-                path = file.path,
-                opts = { visible = false },
-              })
+              if not opts.chunk_mode then
+                -- skip referencing because there will be multiple chunks with the same path (id).
+                -- TODO: figure out a way to deduplicate.
+                agent.chat.references:add({
+                  source = cc_common.tool_result_source,
+                  id = file.path,
+                  path = file.path,
+                  opts = { visible = false },
+                })
+              end
             end
           end
         elseif cmd.command == "ls" then

lua/vectorcode/types.lua

@@ -1,7 +1,10 @@
 ---Type definition of the retrieval result.
 ---@class VectorCode.Result
 ---@field path string Path to the file
----@field document string Content of the file
+---@field document string? Content of the file
+---@field chunk string?
+---@field start_line integer?
+---@field end_line integer?
 
 ---Type definitions for the cache of a buffer.
 ---@class VectorCode.Cache
@@ -53,3 +56,30 @@
 ---@field buf_is_enabled fun(bufnr: integer?): boolean Checks if a buffer has been enabled.
 ---@field make_prompt_component fun(bufnr: integer?, component_cb: (fun(result: VectorCode.Result): string)?): {content: string, count: integer} Compile the retrieval results into a string.
 ---@field async_check fun(check_item: string?, on_success: fun(out: vim.SystemCompleted)?, on_failure: fun(out: vim.SystemCompleted)?) Checks if VectorCode has been configured properly for your project.
+
+--- This class defines the options available to the CodeCompanion tool.
+---@class VectorCode.CodeCompanion.ToolOpts
+--- Maximum number of results provided to the LLM.
+--- You may set this to a table to configure different values for document/chunk mode.
+--- When set to negative values, it means unlimited.
+--- Default: `{ document = -1, chunk = -1 }`
+---@field max_num integer|{document:integer, chunk: integer}|nil
+--- Default number of results provided to the LLM.
+--- This value is written in the system prompt and tool description.
+--- Users may ask the LLM to request a different number of results in the chat.
+--- You may set this to a table to configure different values for document/chunk mode.
+--- Default: `{ document = 10, chunk = 50 }`
+---@field default_num integer|{document:integer, chunk: integer}|nil
+--- Whether the stderr should be provided back to the chat
+---@field include_stderr boolean?
+--- Whether to use the LSP backend. Default: `false`
+---@field use_lsp boolean?
+--- Whether to automatically submit the result (no longer necessary in recent CodeCompanion releases). Default: `false`
+---@field auto_submit table<string, boolean>?
+--- Whether to run `vectorcode ls` and tell the LLM about the indexed projects when initialising the tool. Default: `false`
+---@field ls_on_start boolean?
+--- Whether to avoid duplicated references. Default: `true`
+---@field no_duplicate boolean?
+--- Whether to send chunks instead of full files to the LLM. Default: `false`
+--- > Make sure you adjust `max_num` and `default_num` accordingly.
+---@field chunk_mode boolean?