diff --git a/images/codecompanion_chat.png b/images/codecompanion_chat.png index fb415bc8..c9567e40 100644 Binary files a/images/codecompanion_chat.png and b/images/codecompanion_chat.png differ diff --git a/lua/codecompanion/_extensions/vectorcode/init.lua b/lua/codecompanion/_extensions/vectorcode/init.lua index f6bdf734..dd834b29 100644 --- a/lua/codecompanion/_extensions/vectorcode/init.lua +++ b/lua/codecompanion/_extensions/vectorcode/init.lua @@ -1,72 +1,81 @@ ---@module "codecompanion" +---@alias sub_cmd "ls"|"query"|"vectorise" + ---@class VectorCode.CodeCompanion.ExtensionOpts ----@field add_tools boolean|string|nil ----@field tool_opts VectorCode.CodeCompanion.ToolOpts ----@field add_slash_command boolean +--- A table where the keys are the subcommand name (`ls`, `query`, `vectorise`) +--- and the values are their config options. +---@field tool_opts table +--- Whether to add a tool group that contains all vectorcode tools. +---@field tool_group VectorCode.CodeCompanion.ToolGroupOpts local vc_config = require("vectorcode.config") local logger = vc_config.logger +---@type VectorCode.CodeCompanion.ExtensionOpts|{} +local default_extension_opts = { + tool_opts = { ls = {}, query = {}, vectorise = {} }, + tool_group = { enabled = true, collapse = true, extras = {} }, +} + +---@type sub_cmd[] +local valid_tools = { "ls", "query", "vectorise" } + ---@type CodeCompanion.Extension local M = { ---@param opts VectorCode.CodeCompanion.ExtensionOpts setup = vc_config.check_cli_wrap(function(opts) - opts = vim.tbl_deep_extend( - "force", - { add_tool = true, add_slash_command = false }, - opts or {} - ) + opts = vim.tbl_deep_extend("force", default_extension_opts, opts or {}) logger.info("Received codecompanion extension opts:\n", opts) local cc_config = require("codecompanion.config").config local cc_integration = require("vectorcode.integrations").codecompanion.chat - if opts.add_tool then - local tool_name = "vectorcode" - if type(opts.add_tool) == "string" then - tool_name = tostring(opts.add_tool) - end + for _, sub_cmd in pairs(valid_tools) do + local tool_name = string.format("vectorcode_%s", sub_cmd) if cc_config.strategies.chat.tools[tool_name] ~= nil then vim.notify( - "There's an existing tool named `vectorcode`. Please either remove it or rename it.", + string.format( + "There's an existing tool named `%s`. Please either remove it or rename it.", + tool_name + ), vim.log.levels.ERROR, vc_config.notify_opts ) logger.warn( string.format( - "Not creating a tool because there is an existing tool named %s.", + "Not creating this tool because there is an existing tool named %s.", tool_name ) ) else cc_config.strategies.chat.tools[tool_name] = { - description = "Run VectorCode to retrieve the project context.", - callback = cc_integration.make_tool(opts.tool_opts), + description = string.format("Run VectorCode %s tool", sub_cmd), + callback = cc_integration.make_tool(sub_cmd, opts.tool_opts[sub_cmd]), } logger.info(string.format("%s tool has been created.", tool_name)) end end - if opts.add_slash_command then - local command_name = "vectorcode" - if type(opts.add_slash_command) == "string" then - command_name = tostring(opts.add_slash_command) + + if opts.tool_group.enabled then + local included_tools = vim + .iter(valid_tools) + :map(function(s) + return "vectorcode_" .. s + end) + :totable() + if opts.tool_group.extras and not vim.tbl_isempty(opts.tool_group.extras) then + vim.list_extend(included_tools, opts.tool_group.extras) end - if cc_config.strategies.chat.slash_commands[command_name] ~= nil then - vim.notify( - "There's an existing slash command named `vectorcode`. Please either remove it or rename it.", - vim.log.levels.ERROR, - vc_config.notify_opts - ) - logger.warn( - string.format( - "Not creating a command because there is an existing slash command named %s.", - command_name - ) + logger.info( + string.format( + "Loading the following tools into `vectorcode_toolbox` tool group:\n%s", + vim.inspect(included_tools) ) - else - cc_config.strategies.chat.slash_commands[command_name] = - cc_integration.make_slash_command() - logger.info(string.format("%s command has been created.", command_name)) - end + ) + cc_config.strategies.chat.tools.groups["vectorcode_toolbox"] = { + opts = { collapse_tools = opts.tool_group.collapse }, + description = "Use VectorCode to automatically build and retrieve repository-level context.", + tools = included_tools, + } end end), } diff --git a/lua/vectorcode/cacher/default.lua b/lua/vectorcode/cacher/default.lua index 2b5d9850..b3773fa5 100644 --- a/lua/vectorcode/cacher/default.lua +++ b/lua/vectorcode/cacher/default.lua @@ -257,7 +257,7 @@ M.query_from_cache = vc_config.check_cli_wrap( ---of the array is in the format of `{path="path/to/your/code.lua", document="document content"}`. ---@param bufnr integer? ---@param opts {notify: boolean}? - ---@return VectorCode.Result[] + ---@return VectorCode.QueryResult[] function(bufnr, opts) local result = {} if bufnr == 0 or bufnr == nil then @@ -282,7 +282,7 @@ M.query_from_cache = vc_config.check_cli_wrap( end ) ----@alias ComponentCallback fun(result:VectorCode.Result):string +---@alias ComponentCallback fun(result:VectorCode.QueryResult):string ---Compile the retrieval results into a string. ---@param bufnr integer @@ -296,7 +296,7 @@ function M.make_prompt_component(bufnr, component_cb) return { content = "", count = 0 } end if component_cb == nil then - ---@type fun(result:VectorCode.Result):string + ---@type fun(result:VectorCode.QueryResult):string component_cb = function(result) return "<|file_sep|>" .. result.path .. "\n" .. result.document end diff --git a/lua/vectorcode/cacher/lsp.lua b/lua/vectorcode/cacher/lsp.lua index df9abe68..93e631cb 100644 --- a/lua/vectorcode/cacher/lsp.lua +++ b/lua/vectorcode/cacher/lsp.lua @@ -292,7 +292,7 @@ M.query_from_cache = vc_config.check_cli_wrap( ---of the array is in the format of `{path="path/to/your/code.lua", document="document content"}`. ---@param bufnr integer? ---@param opts {notify: boolean}? - ---@return VectorCode.Result[] + ---@return VectorCode.QueryResult[] function(bufnr, opts) local result = {} if bufnr == 0 or bufnr == nil then @@ -329,7 +329,7 @@ function M.make_prompt_component(bufnr, component_cb) return { content = "", count = 0 } end if component_cb == nil then - ---@type fun(result:VectorCode.Result):string + ---@type fun(result:VectorCode.QueryResult):string component_cb = function(result) return "<|file_sep|>" .. result.path .. "\n" .. result.document end diff --git a/lua/vectorcode/init.lua b/lua/vectorcode/init.lua index ba583889..33bdf96b 100644 --- a/lua/vectorcode/init.lua +++ b/lua/vectorcode/init.lua @@ -14,8 +14,8 @@ M.query = vc_config.check_cli_wrap( ---callback). ---@param query_message string|string[] Query message(s) to send to the `vecctorcode query` command ---@param opts VectorCode.QueryOpts? A table of config options. If nil, the default config or `setup` config will be used. - ---@param callback fun(result:VectorCode.Result[])? Use the result async style. - ---@return VectorCode.Result[]? An array of results. + ---@param callback fun(result:VectorCode.QueryResult[])? Use the result async style. + ---@return VectorCode.QueryResult[]? An array of results. function(query_message, opts, callback) logger.info("vectorcode.query: ", query_message, opts, callback) opts = vim.tbl_deep_extend("force", vc_config.get_query_opts(), opts or {}) diff --git a/lua/vectorcode/integrations/codecompanion/common.lua b/lua/vectorcode/integrations/codecompanion/common.lua index 34e50142..7692d574 100644 --- a/lua/vectorcode/integrations/codecompanion/common.lua +++ b/lua/vectorcode/integrations/codecompanion/common.lua @@ -3,17 +3,22 @@ local vc_config = require("vectorcode.config") local notify_opts = vc_config.notify_opts local logger = vc_config.logger ----@type VectorCode.CodeCompanion.ToolOpts -local default_options = { +---@type VectorCode.CodeCompanion.QueryToolOpts +local default_query_options = { max_num = { chunk = -1, document = -1 }, default_num = { chunk = 50, document = 10 }, - include_stderr = false, use_lsp = false, ls_on_start = false, no_duplicate = true, chunk_mode = false, } +---@type VectorCode.CodeCompanion.LsToolOpts +local default_ls_options = { use_lsp = false } + +---@type VectorCode.CodeCompanion.VectoriseToolOpts +local default_vectorise_options = { use_lsp = false } + return { tool_result_source = "VectorCodeToolResult", ---@param t table|string @@ -25,9 +30,35 @@ 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) + ---@param opts VectorCode.CodeCompanion.LsToolOpts|{}|nil + ---@return VectorCode.CodeCompanion.LsToolOpts + get_ls_tool_opts = function(opts) + opts = vim.tbl_deep_extend("force", default_ls_options, opts or {}) + logger.info( + string.format( + "Loading `vectorcode_ls` with the following opts:\n%s", + vim.inspect(opts) + ) + ) + return opts + end, + + ---@param opts VectorCode.CodeCompanion.VectoriseToolOpts|{}|nil + ---@return VectorCode.CodeCompanion.VectoriseToolOpts + get_vectorise_tool_opts = function(opts) + opts = vim.tbl_deep_extend("force", default_vectorise_options, opts or {}) + logger.info( + string.format( + "Loading `vectorcode_vectorise` with the following opts:\n%s", + vim.inspect(opts) + ) + ) + return opts + end, + + ---@param opts VectorCode.CodeCompanion.QueryToolOpts|{}|nil + ---@return VectorCode.CodeCompanion.QueryToolOpts + get_query_tool_opts = function(opts) if opts == nil or opts.use_lsp == nil then opts = vim.tbl_deep_extend( "force", @@ -35,7 +66,7 @@ return { { use_lsp = vc_config.get_user_config().async_backend == "lsp" } ) end - opts = vim.tbl_deep_extend("force", default_options, opts) + opts = vim.tbl_deep_extend("force", default_query_options, opts) if type(opts.default_num) == "table" then if opts.chunk_mode then opts.default_num = opts.default_num.chunk @@ -48,7 +79,7 @@ return { ) end if type(opts.max_num) == "table" then - if opts.chunk_mode then + if opts._ then opts.max_num = opts.max_num.chunk else opts.max_num = opts.max_num.document @@ -58,10 +89,16 @@ return { "max_num should be an integer or a table: {chunk: integer, document: integer}" ) end + logger.info( + string.format( + "Loading `vectorcode_query` with the following opts:\n%s", + vim.inspect(opts) + ) + ) return opts end, - ---@param result VectorCode.Result + ---@param result VectorCode.QueryResult ---@return string process_result = function(result) local llm_message diff --git a/lua/vectorcode/integrations/codecompanion/func_calling_tool.lua b/lua/vectorcode/integrations/codecompanion/func_calling_tool.lua deleted file mode 100644 index f6963ac0..00000000 --- a/lua/vectorcode/integrations/codecompanion/func_calling_tool.lua +++ /dev/null @@ -1,330 +0,0 @@ ----@module "codecompanion" - -local cc_common = require("vectorcode.integrations.codecompanion.common") -local vc_config = require("vectorcode.config") -local check_cli_wrap = vc_config.check_cli_wrap -local logger = vc_config.logger - -local job_runner = nil - ----@param opts VectorCode.CodeCompanion.ToolOpts? ----@return CodeCompanion.Agent.Tool -return check_cli_wrap(function(opts) - 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 - - logger.info("Creating CodeCompanion tool with the following args:\n", opts) - local capping_message = "" - if opts.max_num > 0 then - capping_message = (" - Request for at most %d documents"):format(opts.max_num) - end - - return { - name = "vectorcode", - cmds = { - ---@param agent CodeCompanion.Agent - ---@param action table - ---@return nil|{ status: string, msg: string } - function(agent, action, _, cb) - logger.info("CodeCompanion tool called with the following arguments:\n", action) - job_runner = cc_common.initialise_runner(opts.use_lsp) - assert(job_runner ~= nil, "Jobrunner not initialised!") - assert( - type(cb) == "function", - "Please upgrade CodeCompanion.nvim to at least 13.5.0" - ) - if not (vim.list_contains({ "ls", "query" }, action.command)) then - if action.options.query ~= nil then - action.command = "query" - else - return { - status = "error", - data = "Need to specify the command (`ls` or `query`).", - } - end - end - - if action.command == "query" then - if action.options.query == nil then - return { - status = "error", - data = "Missing argument: option.query, please refine the tool argument.", - } - end - 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 - if action.options.project_root ~= nil then - action.options.project_root = vim.fs.normalize(action.options.project_root) - if - vim.uv.fs_stat(action.options.project_root) ~= nil - and vim.uv.fs_stat(action.options.project_root).type == "directory" - then - vim.list_extend(args, { "--project_root", action.options.project_root }) - else - return { - status = "error", - data = "INVALID PROJECT ROOT! USE THE LS COMMAND!", - } - end - end - - if opts.no_duplicate and agent.chat.refs ~= nil then - -- exclude files that has been added to the context - local existing_files = { "--exclude" } - for _, ref in pairs(agent.chat.refs) do - if ref.source == cc_common.tool_result_source then - table.insert(existing_files, ref.id) - elseif type(ref.path) == "string" then - table.insert(existing_files, ref.path) - elseif ref.bufnr then - local fname = vim.api.nvim_buf_get_name(ref.bufnr) - if fname ~= nil then - local stat = vim.uv.fs_stat(fname) - if stat and stat.type == "file" then - table.insert(existing_files, fname) - end - end - end - end - if #existing_files > 1 then - vim.list_extend(args, existing_files) - end - end - vim.list_extend(args, { "--absolute" }) - logger.info( - "CodeCompanion query tool called the runner with the following args: ", - args - ) - - job_runner.run_async(args, function(result, error) - if vim.islist(result) and #result > 0 and result[1].path ~= nil then ---@cast result VectorCode.Result[] - cb({ status = "success", data = result }) - else - if type(error) == "table" then - error = cc_common.flatten_table_to_string(error) - end - cb({ - status = "error", - data = error, - }) - end - end, agent.chat.bufnr) - elseif action.command == "ls" then - job_runner.run_async({ "ls", "--pipe" }, function(result, error) - if vim.islist(result) and #result > 0 then - cb({ status = "success", data = result }) - else - if type(error) == "table" then - error = cc_common.flatten_table_to_string(error) - end - cb({ - status = "error", - data = error, - }) - end - end, agent.chat.bufnr) - end - end, - }, - schema = { - type = "function", - ["function"] = { - name = "vectorcode", - description = "Retrieves code documents using semantic search or lists indexed projects", - parameters = { - type = "object", - properties = { - command = { - type = "string", - enum = { "query", "ls" }, - description = "Action to perform: 'query' for semantic search or 'ls' to list projects", - }, - options = { - type = "object", - properties = { - query = { - type = "array", - items = { type = "string" }, - description = "Query messages used for the search.", - }, - count = { - type = "integer", - description = "Number of documents to retrieve, must be positive", - }, - project_root = { - type = "string", - description = "Project path to search within (must be from 'ls' results). Use empty string for the current project.", - }, - }, - required = { "query", "count", "project_root" }, - additionalProperties = false, - }, - }, - required = { "command", "options" }, - additionalProperties = false, - }, - strict = true, - }, - }, - system_prompt = function() - local guidelines = { - " - The path of a retrieved file will be wrapped in `` and `` tags. Its content will be right after the `` tag, wrapped by `` and `` tags. Do not include the ```` 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 . If the starting and ending line ranges are available, they will be wrapped in and 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", - " - Use the `ls` command to retrieve a list of indexed project and pick one that may be relevant, unless the user explicitly mentioned 'this project' (or in other equivalent expressions)", - " - **The project root option MUST be a valid path on the filesystem. It can only be one of the results from the `ls` command or from user input**", - capping_message, - (" - If the user did not specify how many documents to retrieve, **start with %d documents**"):format( - opts.default_num - ), - " - If you decide to call VectorCode tool, do not start answering the question until you have the results. Provide answers based on the results and let the user decide whether to run the tool again", - } - vim.list_extend( - guidelines, - vim.tbl_map(function(line) - return " - " .. line - end, require("vectorcode").prompts({ "query", "ls" })) - ) - if opts.ls_on_start then - job_runner = cc_common.initialise_runner(opts.use_lsp) - if job_runner ~= nil then - local projects = job_runner.run({ "ls", "--pipe" }, -1, 0) - if vim.islist(projects) and #projects > 0 then - vim.list_extend(guidelines, { - " - The following projects are indexed by VectorCode and are available for you to search in:", - }) - vim.list_extend( - guidelines, - vim.tbl_map(function(s) - return string.format(" - %s", s["project-root"]) - end, projects) - ) - end - end - end - local root = vim.fs.root(0, { ".vectorcode", ".git" }) - if root ~= nil then - vim.list_extend(guidelines, { - string.format( - " - The current working directory is %s. Assume the user query is about this project, unless the user asked otherwise or queries from the current project fails to return useful results.", - root - ), - }) - end - return string.format( - [[### VectorCode, a repository indexing and query tool. - -1. **Purpose**: This gives you the ability to access the repository to find information that you may need to assist the user. - -2. **Key Points**: -%s -]], - table.concat(guidelines, "\n") - ) - end, - output = { - ---@param agent CodeCompanion.Agent - ---@param cmd table - ---@param stderr table|string - error = function(self, agent, cmd, stderr) - logger.error( - ("CodeCompanion tool with command %s thrown with the following error: %s"):format( - vim.inspect(cmd), - vim.inspect(stderr) - ) - ) - stderr = cc_common.flatten_table_to_string(stderr) - agent.chat:add_tool_output( - self, - string.format("**VectorCode Tool**: Failed with error:\n```\n%s\n```", stderr) - ) - end, - ---@param agent CodeCompanion.Agent - ---@param cmd table - ---@param stdout table - success = function(self, agent, cmd, stdout) - stdout = stdout[1] - logger.info( - ("CodeCompanion tool with command %s finished."):format(vim.inspect(cmd)) - ) - local user_message - if cmd.command == "query" then - local max_result = #stdout - if opts.max_num > 0 then - 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 = user_message .. " from " .. cmd.options.project_root - end - user_message = user_message .. "\n" - else - user_message = "" - end - agent.chat:add_tool_output( - self, - cc_common.process_result(file), - user_message - ) - 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 - for i, col in pairs(stdout) do - if i == 1 then - user_message = - string.format("Fetched %s indexed project from VectorCode.", #stdout) - else - user_message = "" - end - agent.chat:add_tool_output( - self, - string.format("%s", col["project-root"]), - user_message - ) - end - end - end, - }, - } -end) diff --git a/lua/vectorcode/integrations/codecompanion/init.lua b/lua/vectorcode/integrations/codecompanion/init.lua index 47059ad7..2c52d69f 100644 --- a/lua/vectorcode/integrations/codecompanion/init.lua +++ b/lua/vectorcode/integrations/codecompanion/init.lua @@ -5,7 +5,7 @@ local check_cli_wrap = vc_config.check_cli_wrap return { chat = { - ---@param component_cb (fun(result:VectorCode.Result):string)? + ---@param component_cb (fun(result:VectorCode.QueryResult):string)? make_slash_command = check_cli_wrap(function(component_cb) return { description = "Add relevant files from the codebase.", @@ -35,12 +35,15 @@ return { } end), + ---@param subcommand "ls"|"query"|"vectorise" ---@param opts VectorCode.CodeCompanion.ToolOpts ---@return CodeCompanion.Agent.Tool - make_tool = function(opts) + make_tool = function(subcommand, opts) local has = require("codecompanion").has if has ~= nil and has("function-calling") then - return require("vectorcode.integrations.codecompanion.func_calling_tool")(opts) + return require( + string.format("vectorcode.integrations.codecompanion.%s_tool", subcommand) + )(opts) else error("Unsupported version of codecompanion!") end diff --git a/lua/vectorcode/integrations/codecompanion/ls_tool.lua b/lua/vectorcode/integrations/codecompanion/ls_tool.lua new file mode 100644 index 00000000..eb0bcc3d --- /dev/null +++ b/lua/vectorcode/integrations/codecompanion/ls_tool.lua @@ -0,0 +1,69 @@ +---@module "codecompanion" + +local cc_common = require("vectorcode.integrations.codecompanion.common") +local vectorcode = require("vectorcode") + +---@param opts VectorCode.CodeCompanion.LsToolOpts +---@return CodeCompanion.Agent.Tool +return function(opts) + opts = cc_common.get_ls_tool_opts(opts) + local job_runner = + require("vectorcode.integrations.codecompanion.common").initialise_runner( + opts.use_lsp + ) + local tool_name = "vectorcode_ls" + ---@type CodeCompanion.Agent.Tool|{} + return { + name = tool_name, + cmds = { + ---@param agent CodeCompanion.Agent + ---@return nil|{ status: string, data: string } + function(agent, _, _, cb) + job_runner.run_async({ "ls", "--pipe" }, function(result, error) + if vim.islist(result) and #result > 0 then + cb({ status = "success", data = result }) + else + if type(error) == "table" then + error = cc_common.flatten_table_to_string(error) + end + cb({ + status = "error", + data = error, + }) + end + end, agent.chat.bufnr) + end, + }, + schema = { + type = "function", + ["function"] = { + name = tool_name, + description = string.format( + "Retrieve a list of projects accessible via the VectorCode tools.\n%s", + table.concat(vectorcode.prompts("ls"), "\n") + ), + }, + }, + output = { + ---@param agent CodeCompanion.Agent + ---@param stdout VectorCode.LsResult[][] + success = function(_, agent, _, stdout) + stdout = stdout[1] + local user_message + for i, col in ipairs(stdout) do + if i == 1 then + user_message = + string.format("**VectorCode `ls` Tool**: Found %d collections.", #stdout) + else + user_message = "" + end + agent.chat:add_tool_output( + agent.tool, + string.format("%s", col["project-root"]), + user_message + ) + end + end, + }, + } +end diff --git a/lua/vectorcode/integrations/codecompanion/query_tool.lua b/lua/vectorcode/integrations/codecompanion/query_tool.lua new file mode 100644 index 00000000..9f893ce9 --- /dev/null +++ b/lua/vectorcode/integrations/codecompanion/query_tool.lua @@ -0,0 +1,237 @@ +---@module "codecompanion" + +local cc_common = require("vectorcode.integrations.codecompanion.common") +local vc_config = require("vectorcode.config") +local check_cli_wrap = vc_config.check_cli_wrap +local logger = vc_config.logger + +local job_runner = nil + +---@alias QueryToolArgs { project_root:string, count: integer, query: string[] } + +---@param opts VectorCode.CodeCompanion.QueryToolOpts? +---@return CodeCompanion.Agent.Tool +return check_cli_wrap(function(opts) + opts = cc_common.get_query_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 + + logger.info("Creating CodeCompanion tool with the following args:\n", opts) + + local tool_name = "vectorcode_query" + return { + name = tool_name, + cmds = { + ---@param agent CodeCompanion.Agent + ---@param action QueryToolArgs + ---@return nil|{ status: string, data: string } + function(agent, action, _, cb) + logger.info( + "CodeCompanion query tool called with the following arguments:\n", + action + ) + job_runner = cc_common.initialise_runner(opts.use_lsp) + assert(job_runner ~= nil, "Jobrunner not initialised!") + assert( + type(cb) == "function", + "Please upgrade CodeCompanion.nvim to at least 13.5.0" + ) + + if action.query == nil then + return { + status = "error", + data = "Missing argument: option.query, please refine the tool argument.", + } + end + + local args = { "query" } + vim.list_extend(args, action.query) + vim.list_extend(args, { "--pipe", "-n", tostring(action.count) }) + if opts.chunk_mode then + vim.list_extend(args, { "--include", "path", "chunk" }) + else + vim.list_extend(args, { "--include", "path", "document" }) + end + if action.project_root == "" then + action.project_root = nil + end + if action.project_root ~= nil then + action.project_root = vim.fs.normalize(action.project_root) + if + vim.uv.fs_stat(action.project_root) ~= nil + and vim.uv.fs_stat(action.project_root).type == "directory" + then + action.project_root = vim.fs.abspath(vim.fs.normalize(action.project_root)) + vim.list_extend(args, { "--project_root", action.project_root }) + else + return { + status = "error", + data = "INVALID PROJECT ROOT! USE THE LS COMMAND!", + } + end + end + + if opts.no_duplicate and agent.chat.refs ~= nil then + -- exclude files that has been added to the context + local existing_files = { "--exclude" } + for _, ref in pairs(agent.chat.refs) do + if ref.source == cc_common.tool_result_source then + table.insert(existing_files, ref.id) + elseif type(ref.path) == "string" then + table.insert(existing_files, ref.path) + elseif ref.bufnr then + local fname = vim.api.nvim_buf_get_name(ref.bufnr) + if fname ~= nil then + local stat = vim.uv.fs_stat(fname) + if stat and stat.type == "file" then + table.insert(existing_files, fname) + end + end + end + end + if #existing_files > 1 then + vim.list_extend(args, existing_files) + end + end + vim.list_extend(args, { "--absolute" }) + logger.info( + "CodeCompanion query tool called the runner with the following args: ", + args + ) + + job_runner.run_async(args, function(result, error) + if vim.islist(result) and #result > 0 and result[1].path ~= nil then ---@cast result VectorCode.QueryResult[] + cb({ status = "success", data = result }) + else + if type(error) == "table" then + error = cc_common.flatten_table_to_string(error) + end + cb({ + status = "error", + data = error, + }) + end + end, agent.chat.bufnr) + end, + }, + schema = { + type = "function", + ["function"] = { + name = tool_name, + description = [[Retrieves code documents using semantic search. +The path of a retrieved file will be wrapped in `` and `` tags. +Its content will be right after the `` tag, wrapped by `` and `` tags. +Do not include the xml 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 . +If the starting and ending line ranges are available, they will be wrapped in and tags. +Make use of the line numbers (NOT THE XML TAGS) when you're quoting the source code. +Include one single command call for VectorCode each time. +You may include multiple keywords in the command. +**The project root option MUST be a valid path on the filesystem. It can only be one of the results from the `vectorcode_ls` tool or from user input** + ]], + parameters = { + type = "object", + properties = { + query = { + type = "array", + items = { type = "string" }, + description = [[ +Query messages used for the search. They should also contain relevant keywords. +For example, you should include `parameter`, `arguments` and `return value` for the query `function`. + ]], + }, + count = { + type = "integer", + description = string.format( + "Number of documents or chunks to retrieve, must be positive. Use %d by default. Do not query for more than %d.", + tonumber(opts.default_num), + tonumber(opts.max_num) + ), + }, + project_root = { + type = "string", + description = "Project path to search within (must be from 'ls' results or user instructions). Use empty string for the current project.", + }, + }, + required = { "query", "count", "project_root" }, + additionalProperties = false, + }, + strict = true, + }, + }, + output = { + ---@param agent CodeCompanion.Agent + ---@param cmd QueryToolArgs + ---@param stderr table|string + error = function(self, agent, cmd, stderr) + logger.error( + ("CodeCompanion tool with command %s thrown with the following error: %s"):format( + vim.inspect(cmd), + vim.inspect(stderr) + ) + ) + stderr = cc_common.flatten_table_to_string(stderr) + agent.chat:add_tool_output( + self, + string.format("**VectorCode Tool**: Failed with error:\n```\n%s\n```", stderr) + ) + end, + ---@param agent CodeCompanion.Agent + ---@param cmd QueryToolArgs + ---@param stdout VectorCode.QueryResult[][] + success = function(self, agent, cmd, stdout) + stdout = stdout[1] + logger.info( + ("CodeCompanion tool with command %s finished."):format(vim.inspect(cmd)) + ) + local user_message + local max_result = #stdout + if opts.max_num > 0 then + 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.project_root then + user_message = user_message .. " from " .. cmd.project_root + end + user_message = user_message .. "\n" + else + user_message = "" + end + agent.chat:add_tool_output( + self, + cc_common.process_result(file), + user_message + ) + 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 + end, + }, + } +end) diff --git a/lua/vectorcode/integrations/codecompanion/vectorise_tool.lua b/lua/vectorcode/integrations/codecompanion/vectorise_tool.lua new file mode 100644 index 00000000..7cb15953 --- /dev/null +++ b/lua/vectorcode/integrations/codecompanion/vectorise_tool.lua @@ -0,0 +1,127 @@ +---@module "codecompanion" + +local cc_common = require("vectorcode.integrations.codecompanion.common") +local vectorcode = require("vectorcode") + +---@alias VectoriseToolArgs { paths: string[], project_root: string } + +---@alias VectoriseResult { add: integer, update: integer, removed: integer } + +---@param opts VectorCode.CodeCompanion.VectoriseToolOpts|{}|nil +---@return CodeCompanion.Agent.Tool +return function(opts) + opts = cc_common.get_vectorise_tool_opts(opts) + local tool_name = "vectorcode_vectorise" + local job_runner = cc_common.initialise_runner(opts.use_lsp) + + ---@type CodeCompanion.Agent.Tool|{} + return { + name = tool_name, + schema = { + type = "function", + ["function"] = { + name = tool_name, + description = string.format( + "Vectorise files in a project so that they'll be available from the vectorcode_query tool\n%s", + table.concat(vectorcode.prompts("vectorise"), "\n") + ), + parameters = { + type = "object", + properties = { + paths = { + type = "array", + items = { type = "string" }, + description = "Paths to the files to be vectorised", + }, + project_root = { + type = "string", + description = "The project that the files belong to. Either use a path from the `vectorcode_ls` tool, or leave empty to use the current git project.", + }, + }, + required = { "paths", "project_root" }, + additionalProperties = false, + }, + strict = true, + }, + }, + cmds = { + ---@param agent CodeCompanion.Agent + ---@param action VectoriseToolArgs + ---@return nil|{ status: string, data: string } + function(agent, action, _, cb) + local args = { "vectorise", "--pipe" } + local project_root = vim.fs.abspath(vim.fs.normalize(action.project_root or "")) + if project_root ~= "" then + local stat = vim.uv.fs_stat(project_root) + if stat and stat.type == "directory" then + vim.list_extend(args, { "--project_root", project_root }) + else + return { status = "error", data = "Invalid path " .. project_root } + end + else + project_root = vim.fs.root(".", { ".vectorcode", ".git" }) or "" + if project_root == "" then + return { + status = "error", + data = "Please specify a project root. You may use the `vectorcode_ls` tool to find a list of existing projects.", + } + end + end + vim.list_extend( + args, + vim + .iter(action.paths) + :filter( + ---@param item string + function(item) + local stat = vim.uv.fs_stat(item) + if stat and stat.type == "file" then + return true + else + return false + end + end + ) + :totable() + ) + job_runner.run_async( + args, + ---@param result VectoriseResult + function(result, error, code, _) + if result then + cb({ status = "success", data = result }) + else + cb({ status = "error", data = { error = error, code = code } }) + end + end, + agent.chat.bufnr + ) + end, + }, + output = { + ---@param self CodeCompanion.Agent.Tool + ---@param agent CodeCompanion.Agent + ---@param stdout VectorCode.VectoriseResult[] + success = function(self, agent, _, stdout) + stdout = stdout[1] + agent.chat:add_tool_output( + self, + string.format( + [[**VectorCode Vectorise Tool**: + - New files added: %d + - Existing files updated: %d + - Orphaned files removed: %d + - Up-to-date files skipped: %d + - Failed to decode: %d + ]], + stdout.add, + stdout.update, + stdout.removed, + stdout.skipped, + stdout.failed + ) + ) + end, + }, + } +end diff --git a/lua/vectorcode/types.lua b/lua/vectorcode/types.lua index a174c542..c53c62d4 100644 --- a/lua/vectorcode/types.lua +++ b/lua/vectorcode/types.lua @@ -1,11 +1,21 @@ ---Type definition of the retrieval result. ----@class VectorCode.Result +---@class VectorCode.QueryResult ---@field path string Path to the file ---@field document string? Content of the file ---@field chunk string? ---@field start_line integer? ---@field end_line integer? +---@class VectorCode.LsResult +---@field project-root string + +---@class VectorCode.VectoriseResult +---@field add integer +---@field update integer +---@field removed integer +---@field skipped integer +---@field failed integer + ---Type definitions for the cache of a buffer. ---@class VectorCode.Cache ---@field enabled boolean Whether the async jobs are enabled or not. If the buffer is disabled, no cache will be generated for it. @@ -13,7 +23,7 @@ ---@field jobs table Job handle:time of creation (in seconds) ---@field last_run integer? Last time the query ran, in seconds from epoch. ---@field options VectorCode.RegisterOpts The options that the buffer was registered with. ----@field retrieval VectorCode.Result[]? The latest retrieval. +---@field retrieval VectorCode.QueryResult[]? The latest retrieval. ---Type definitions for options accepted by `query` API. ---@class VectorCode.QueryOpts @@ -50,15 +60,21 @@ ---@class VectorCode.CacheBackend ---@field register_buffer fun(bufnr: integer?, opts: VectorCode.RegisterOpts) Register a buffer and create an async cache for it. ---@field deregister_buffer fun(bufnr: integer?, opts: {notify: boolean}?) Deregister a buffer and destroy its async cache. ----@field query_from_cache fun(bufnr: integer?, opts: {notify: boolean}?): VectorCode.Result[] Get the cached documents. +---@field query_from_cache fun(bufnr: integer?, opts: {notify: boolean}?): VectorCode.QueryResult[] Get the cached documents. ---@field buf_is_registered fun(bufnr: integer?): boolean Checks if a buffer has been registered. ---@field buf_job_count fun(bufnr: integer?): integer Returns the number of running jobs in the background. ---@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 make_prompt_component fun(bufnr: integer?, component_cb: (fun(result: VectorCode.QueryResult): 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 +--- Whether to use the LSP backend. Default: `false` +---@field use_lsp boolean? + +---@class VectorCode.CodeCompanion.LsToolOpts: VectorCode.CodeCompanion.ToolOpts + +---@class VectorCode.CodeCompanion.QueryToolOpts: 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. @@ -70,16 +86,18 @@ --- 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? ---- 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? + +---@class VectorCode.CodeCompanion.VectoriseToolOpts: VectorCode.CodeCompanion.ToolOpts + +---@class VectorCode.CodeCompanion.ToolGroupOpts +--- Whether to register the tool group +---@field enabled boolean +--- Whether to show the individual tools in the references +---@field collapse boolean +--- Other tools that you'd like to include in `vectorcode_toolbox` +---@field extras string[]