From fe6518f6de91853f6249a3790cf5aebfcc514b41 Mon Sep 17 00:00:00 2001 From: yetone Date: Tue, 27 Aug 2024 22:44:40 +0800 Subject: [PATCH] feat: editing mode (#281) --- README.md | 2 +- lua/avante/init.lua | 21 ++- lua/avante/llm.lua | 24 ++- lua/avante/selection.lua | 371 ++++++++++++++++++++++++++++++++++++-- lua/avante/sidebar.lua | 5 + lua/avante/utils/init.lua | 6 + 6 files changed, 404 insertions(+), 25 deletions(-) diff --git a/README.md b/README.md index d9945d4..ace7683 100644 --- a/README.md +++ b/README.md @@ -214,7 +214,7 @@ The following key bindings are available for use with `avante.nvim`: - [x] Apply diff patch - [x] Chat with the selected block - [x] Slash commands -- [ ] Edit the selected block +- [x] Edit the selected block - [ ] Smart Tab (Cursor Flow) - [ ] Chat with project - [ ] Chat with selected files diff --git a/lua/avante/init.lua b/lua/avante/init.lua index cc96b30..35313b4 100644 --- a/lua/avante/init.lua +++ b/lua/avante/init.lua @@ -29,7 +29,7 @@ H.commands = function() M.toggle() end, { desc = "avante: ask AI for code suggestions" }) cmd("Close", function() - local sidebar = M._get() + local sidebar, _ = M._get() if not sidebar then return end @@ -42,6 +42,7 @@ end H.keymaps = function() vim.keymap.set({ "n", "v" }, Config.mappings.ask, M.toggle, { noremap = true }) + vim.keymap.set("v", Config.mappings.edit, M.edit, { noremap = true }) vim.keymap.set("n", Config.mappings.refresh, M.refresh, { noremap = true }) Utils.toggle_map("n", Config.mappings.toggle.debug, { @@ -116,7 +117,7 @@ H.autocmds = function() api.nvim_create_autocmd("VimResized", { group = H.augroup, callback = function() - local sidebar = M._get() + local sidebar, _ = M._get() if not sidebar then return end @@ -165,7 +166,7 @@ H.autocmds = function() end ---@param current boolean? false to disable setting current, otherwise use this to track across tabs. ----@return avante.Sidebar +---@return avante.Sidebar, avante.Selection function M._get(current) local tab = api.nvim_get_current_tabpage() local sidebar = M.sidebars[tab] @@ -174,7 +175,7 @@ function M._get(current) M.current.sidebar = sidebar M.current.selection = selection end - return sidebar + return sidebar, selection end ---@param id integer @@ -195,7 +196,7 @@ function M._init(id) end M.toggle = function() - local sidebar = M._get() + local sidebar, _ = M._get() if not sidebar then M._init(api.nvim_get_current_tabpage()) M.current.sidebar:open() @@ -205,8 +206,16 @@ M.toggle = function() return sidebar:toggle() end +M.edit = function() + local _, selection = M._get() + if not selection then + return + end + selection:create_editing_input() +end + M.refresh = function() - local sidebar = M._get() + local sidebar, _ = M._get() if not sidebar then return end diff --git a/lua/avante/llm.lua b/lua/avante/llm.lua index 3962db4..49e79b0 100644 --- a/lua/avante/llm.lua +++ b/lua/avante/llm.lua @@ -19,7 +19,7 @@ You are an excellent programming expert. ]] ---@alias AvanteBasePrompt string -local base_user_prompt = [[ +local planning_mode_prompt = [[ Your primary task is to suggest code modifications with precise line number ranges. Follow these instructions meticulously: 1. Carefully analyze the original code, paying close attention to its structure and line numbers. Line numbers start from 1 and include ALL lines, even empty ones. @@ -44,6 +44,7 @@ Replace lines: {{start_line}}-{{end_line}} - Do not omit any needed changes from the requisite messages/code blocks. - If there is a clicked code block, bias towards just applying that (and applying other changes implied). - Please keep your suggested code changes minimal, and do not include irrelevant lines in the code snippet. + - Maintain the SAME indentation in the returned code as in the source code 4. Crucial guidelines for line numbers: - The content regarding line numbers MUST strictly follow the format "Replace lines: {{start_line}}-{{end_line}}". Do not be lazy! @@ -64,6 +65,21 @@ Replace lines: {{start_line}}-{{end_line}} Remember: Accurate line numbers are CRITICAL. The range start_line to end_line must include ALL lines to be replaced, from the very first to the very last. Double-check every range before finalizing your response, paying special attention to the start_line to ensure it hasn't shifted down. Ensure that your line numbers perfectly match the original code structure without any overall shift. ]] +local editing_mode_prompt = [[ +Your task is to modify the provided code according to the user's request. Follow these instructions precisely: + +1. Carefully analyze the original code and the user's request. +2. Make the necessary modifications to the code as requested. +3. Return ONLY the complete modified code. +4. Do not include any explanations, comments, or line numbers in your response. +5. Ensure the returned code is complete and can be directly used as a replacement for the original code. +6. Preserve the original structure, indentation, and formatting of the code as much as possible. +7. Do not omit any parts of the code, even if they are unchanged. +8. Maintain the SAME indentation in the returned code as in the source code + +Remember: Your response should contain nothing but the modified code, ready to be used as a direct replacement for the original file. +]] + local group = api.nvim_create_augroup("avante_llm", { clear = true }) local active_job = nil @@ -71,14 +87,16 @@ local active_job = nil ---@param code_lang string ---@param code_content string ---@param selected_content_content string | nil +---@param mode "planning" | "editing" ---@param on_chunk AvanteChunkParser ---@param on_complete AvanteCompleteParser -M.stream = function(question, code_lang, code_content, selected_content_content, on_chunk, on_complete) +M.stream = function(question, code_lang, code_content, selected_content_content, mode, on_chunk, on_complete) + mode = mode or "planning" local provider = Config.provider ---@type AvantePromptOptions local code_opts = { - base_prompt = base_user_prompt, + base_prompt = mode == "planning" and planning_mode_prompt or editing_mode_prompt, system_prompt = system_prompt, question = question, code_lang = code_lang, diff --git a/lua/avante/selection.lua b/lua/avante/selection.lua index 1925aa6..0d64c5b 100644 --- a/lua/avante/selection.lua +++ b/lua/avante/selection.lua @@ -1,5 +1,7 @@ local Utils = require("avante.utils") local Config = require("avante.config") +local Llm = require("avante.llm") +local Highlights = require("avante.highlights") local api = vim.api local fn = vim.fn @@ -7,7 +9,18 @@ local fn = vim.fn local NAMESPACE = api.nvim_create_namespace("avante_selection") local PRIORITY = vim.highlight.priorities.user +local EIDTING_INPUT_START_SPINNER_PATTERN = "AvanteEditingInputStartSpinner" +local EIDTING_INPUT_STOP_SPINNER_PATTERN = "AvanteEditingInputStopSpinner" + ---@class avante.Selection +---@field selection avante.SelectionResult | nil +---@field cursor_pos table | nil +---@field shortcuts_extmark_id integer | nil +---@field augroup integer | nil +---@field editing_input_bufnr integer | nil +---@field editing_input_winid integer | nil +---@field editing_input_shortcuts_hints_winid integer | nil +---@field code_winid integer | nil local Selection = {} Selection.did_setup = false @@ -15,9 +28,14 @@ Selection.did_setup = false ---@param id integer the tabpage id retrieved from api.nvim_get_current_tabpage() function Selection:new(id) return setmetatable({ - hints_popup_extmark_id = nil, - edit_popup_renderer = nil, + shortcuts_extmark_id = nil, augroup = api.nvim_create_augroup("avante_selection_" .. id, { clear = true }), + selection = nil, + cursor_pos = nil, + editing_input_bufnr = nil, + editing_input_winid = nil, + editing_input_shortcuts_hints_winid = nil, + code_winid = nil, }, { __index = self }) end @@ -40,27 +58,350 @@ function Selection:get_virt_text_line() return current_line end -function Selection:show_hints_popup() - self:close_hints_popup() +function Selection:show_shortcuts_hints_popup() + self:close_shortcuts_hints_popup() - local hint_text = string.format(" [%s: ask avante] ", Config.mappings.ask) + local hint_text = string.format(" [%s: ask avante, %s: edit] ", Config.mappings.ask, Config.mappings.edit) local virt_text_line = self:get_virt_text_line() - self.hints_popup_extmark_id = api.nvim_buf_set_extmark(0, NAMESPACE, virt_text_line, -1, { + self.shortcuts_extmark_id = api.nvim_buf_set_extmark(0, NAMESPACE, virt_text_line, -1, { virt_text = { { hint_text, "Keyword" } }, virt_text_pos = "eol", priority = PRIORITY, }) end -function Selection:close_hints_popup() - if self.hints_popup_extmark_id then - api.nvim_buf_del_extmark(0, NAMESPACE, self.hints_popup_extmark_id) - self.hints_popup_extmark_id = nil +function Selection:close_shortcuts_hints_popup() + if self.shortcuts_extmark_id then + api.nvim_buf_del_extmark(0, NAMESPACE, self.shortcuts_extmark_id) + self.shortcuts_extmark_id = nil end end +function Selection:close_editing_input() + self:close_editing_input_shortcuts_hints() + api.nvim_exec_autocmds("User", { pattern = Llm.CANCEL_PATTERN }) + if api.nvim_get_mode().mode == "i" then + vim.cmd([[stopinsert]]) + end + if self.editing_input_winid and api.nvim_win_is_valid(self.editing_input_winid) then + api.nvim_win_close(self.editing_input_winid, true) + self.editing_input_winid = nil + end + if self.cursor_pos and self.code_winid then + vim.schedule(function() + api.nvim_win_set_cursor(self.code_winid, { self.cursor_pos[1], self.cursor_pos[2] }) + end) + end + if self.editing_input_bufnr and api.nvim_buf_is_valid(self.editing_input_bufnr) then + api.nvim_buf_delete(self.editing_input_bufnr, { force = true }) + self.editing_input_bufnr = nil + end +end + +function Selection:close_editing_input_shortcuts_hints() + if self.editing_input_shortcuts_hints_winid and api.nvim_win_is_valid(self.editing_input_shortcuts_hints_winid) then + api.nvim_win_close(self.editing_input_shortcuts_hints_winid, true) + self.editing_input_shortcuts_hints_winid = nil + end +end + +function Selection:show_editing_input_shortcuts_hints() + self:close_editing_input_shortcuts_hints() + + if not self.editing_input_winid or not api.nvim_win_is_valid(self.editing_input_winid) then + return + end + + local win_width = api.nvim_win_get_width(self.editing_input_winid) + local buf_height = api.nvim_buf_line_count(self.editing_input_bufnr) + -- "⠁⠁⠉⠙⠚⠒⠂⠂⠒⠲⠴⠤⠄⠄⠤⢤⣠⡀⡀⣀⢀⢀⣄⡤⠤⠠⠠⠤⠦⠖⠒⠐⠐⠒⠓⠋⠉⠈⠈⠉" + local spinner_chars = { + "⠁", + "⠉", + "⠙", + "⠚", + "⠒", + "⠂", + "⠂", + "⠒", + "⠲", + "⠴", + "⠤", + "⠄", + "⠄", + "⠤", + "⢤", + "⣠", + "⡀", + "⡀", + "⣀", + "⢀", + "⢀", + "⣄", + "⡤", + "⠤", + "⠠", + "⠠", + "⠤", + "⠦", + "⠖", + "⠒", + "⠐", + "⠐", + "⠒", + "⠓", + "⠋", + "⠉", + "⠈", + "⠈", + "⠉", + } + local spinner_index = 1 + local timer = nil + + local hint_text = (vim.fn.mode() ~= "i" and Config.mappings.submit.normal or Config.mappings.submit.insert) + .. ": submit" + + local buf = api.nvim_create_buf(false, true) + api.nvim_buf_set_lines(buf, 0, -1, false, { hint_text }) + + local function update_spinner() + spinner_index = (spinner_index % #spinner_chars) + 1 + local spinner = spinner_chars[spinner_index] + local new_text = spinner .. " " .. hint_text + + api.nvim_buf_set_lines(buf, 0, -1, false, { new_text }) + + if + not self.editing_input_shortcuts_hints_winid + or not api.nvim_win_is_valid(self.editing_input_shortcuts_hints_winid) + then + return + end + + local win_config = vim.api.nvim_win_get_config(self.editing_input_shortcuts_hints_winid) + + if win_config.width ~= #new_text then + win_config.width = #new_text + win_config.col = math.max(win_width - #new_text, 0) + vim.api.nvim_win_set_config(self.editing_input_shortcuts_hints_winid, win_config) + end + end + + local function stop_spinner() + if timer then + timer:stop() + timer:close() + timer = nil + end + api.nvim_buf_set_lines(buf, 0, -1, false, { hint_text }) + + if + not self.editing_input_shortcuts_hints_winid + or not api.nvim_win_is_valid(self.editing_input_shortcuts_hints_winid) + then + return + end + + local win_config = vim.api.nvim_win_get_config(self.editing_input_shortcuts_hints_winid) + + if win_config.width ~= #hint_text then + win_config.width = #hint_text + win_config.col = math.max(win_width - #hint_text, 0) + vim.api.nvim_win_set_config(self.editing_input_shortcuts_hints_winid, win_config) + end + end + + api.nvim_create_autocmd("User", { + pattern = EIDTING_INPUT_START_SPINNER_PATTERN, + callback = function() + timer = vim.loop.new_timer() + if timer then + timer:start( + 0, + 100, + vim.schedule_wrap(function() + update_spinner() + end) + ) + end + end, + }) + + api.nvim_create_autocmd("User", { + pattern = EIDTING_INPUT_STOP_SPINNER_PATTERN, + callback = function() + stop_spinner() + end, + }) + local width = #hint_text + + local opts = { + relative = "win", + win = self.editing_input_winid, + width = width, + height = 1, + row = buf_height, + col = math.max(win_width - width, 0), + style = "minimal", + border = "none", + focusable = false, + zindex = 100, + } + + self.editing_input_shortcuts_hints_winid = api.nvim_open_win(buf, false, opts) + + api.nvim_win_set_hl_ns(self.editing_input_shortcuts_hints_winid, Highlights.hint_ns) +end + +function Selection:create_editing_input() + local code_bufnr = api.nvim_get_current_buf() + local code_wind = api.nvim_get_current_win() + self.cursor_pos = api.nvim_win_get_cursor(code_wind) + self.code_winid = code_wind + local filetype = api.nvim_get_option_value("filetype", { buf = code_bufnr }) + local code_lines = api.nvim_buf_get_lines(code_bufnr, 0, -1, false) + local code_content = table.concat(code_lines, "\n") + + self.selection = Utils.get_visual_selection_and_range() + + self:close_editing_input() + + local bufnr = api.nvim_create_buf(false, true) + + self.editing_input_bufnr = bufnr + + local win_opts = { + relative = "cursor", + width = 40, + height = 2, + row = 1, + col = 0, + style = "minimal", + border = "rounded", + title = { { "Chat with selected code", "FloatTitle" } }, + title_pos = "center", + } + + local winid = api.nvim_open_win(bufnr, true, win_opts) + + self.editing_input_winid = winid + + api.nvim_set_option_value("wrap", false, { win = winid }) + api.nvim_set_option_value("cursorline", true, { win = winid }) + api.nvim_set_option_value("modifiable", true, { buf = bufnr }) + + self:show_editing_input_shortcuts_hints() + + api.nvim_create_autocmd({ "TextChanged", "TextChangedI" }, { + group = self.augroup, + buffer = bufnr, + callback = function() + self:show_editing_input_shortcuts_hints() + end, + }) + + api.nvim_create_autocmd("ModeChanged", { + group = self.augroup, + pattern = "i:*", + callback = function() + local cur_buf = api.nvim_get_current_buf() + if cur_buf == bufnr then + self:show_editing_input_shortcuts_hints() + end + end, + }) + + api.nvim_create_autocmd("ModeChanged", { + group = self.augroup, + pattern = "*:i", + callback = function() + local cur_buf = api.nvim_get_current_buf() + if cur_buf == bufnr then + self:show_editing_input_shortcuts_hints() + end + end, + }) + + local function submit_input() + local lines = api.nvim_buf_get_lines(bufnr, 0, -1, false) + local input = lines[1] or "" + + local full_response = "" + local start_line = self.selection.range.start.line + local finish_line = self.selection.range.finish.line + + local indentation = Utils.get_indentation(code_lines[self.selection.range.start.line]) + + api.nvim_exec_autocmds("User", { pattern = EIDTING_INPUT_START_SPINNER_PATTERN }) + ---@type AvanteChunkParser + local on_chunk = function(chunk) + full_response = full_response .. chunk + local response_lines = vim.split(full_response, "\n") + for i, line in ipairs(response_lines) do + response_lines[i] = indentation .. line + end + api.nvim_buf_set_lines(code_bufnr, start_line - 1, finish_line, true, response_lines) + finish_line = start_line + #response_lines - 1 + end + + ---@type AvanteCompleteParser + local on_complete = function(err) + if err then + Utils.error( + "Error occurred while processing the response: " .. vim.inspect(err), + { once = true, title = "Avante" } + ) + return + end + api.nvim_exec_autocmds("User", { pattern = EIDTING_INPUT_STOP_SPINNER_PATTERN }) + vim.defer_fn(function() + self:close_editing_input() + end, 0) + end + + Llm.stream(input, filetype, code_content, self.selection.content, "editing", on_chunk, on_complete) + end + + vim.keymap.set("i", Config.mappings.submit.insert, submit_input, { buffer = bufnr, noremap = true, silent = true }) + vim.keymap.set("n", Config.mappings.submit.normal, submit_input, { buffer = bufnr, noremap = true, silent = true }) + vim.keymap.set("n", "", function() + self:close_editing_input() + end, { buffer = bufnr }) + vim.keymap.set("n", "q", function() + self:close_editing_input() + end, { buffer = bufnr }) + + local quit_id, close_unfocus + quit_id = api.nvim_create_autocmd("QuitPre", { + group = self.augroup, + buffer = bufnr, + once = true, + nested = true, + callback = function() + self:close_editing_input() + if not quit_id then + api.nvim_del_autocmd(quit_id) + quit_id = nil + end + end, + }) + + close_unfocus = api.nvim_create_autocmd("WinLeave", { + group = self.augroup, + buffer = bufnr, + callback = function() + self:close_editing_input() + if close_unfocus then + api.nvim_del_autocmd(close_unfocus) + close_unfocus = nil + end + end, + }) +end + function Selection:setup_autocmds() Selection.did_setup = true api.nvim_create_autocmd({ "ModeChanged" }, { @@ -68,7 +409,7 @@ function Selection:setup_autocmds() pattern = { "n:v", "n:V", "n:" }, -- Entering Visual mode from Normal mode callback = function(ev) if not Utils.is_sidebar_buffer(ev.buf) then - self:show_hints_popup() + self:show_shortcuts_hints_popup() end end, }) @@ -78,9 +419,9 @@ function Selection:setup_autocmds() callback = function(ev) if not Utils.is_sidebar_buffer(ev.buf) then if Utils.in_visual_mode() then - self:show_hints_popup() + self:show_shortcuts_hints_popup() else - self:close_hints_popup() + self:close_shortcuts_hints_popup() end end end, @@ -91,7 +432,7 @@ function Selection:setup_autocmds() pattern = { "v:n", "v:i", "v:c" }, -- Switching from visual mode back to normal, insert, or other modes callback = function(ev) if not Utils.is_sidebar_buffer(ev.buf) then - self:close_hints_popup() + self:close_shortcuts_hints_popup() end end, }) @@ -100,7 +441,7 @@ function Selection:setup_autocmds() group = self.augroup, callback = function(ev) if not Utils.is_sidebar_buffer(ev.buf) then - self:close_hints_popup() + self:close_shortcuts_hints_popup() end end, }) diff --git a/lua/avante/sidebar.lua b/lua/avante/sidebar.lua index 42da50b..286306a 100644 --- a/lua/avante/sidebar.lua +++ b/lua/avante/sidebar.lua @@ -1008,6 +1008,10 @@ function Sidebar:create_input() self.input:unmount() end + if not self.code.bufnr or not api.nvim_buf_is_valid(self.code.bufnr) then + return + end + local chat_history = History.load(self.code.bufnr) ---@param request string @@ -1133,6 +1137,7 @@ function Sidebar:create_input() filetype, content_with_line_numbers, selected_code_content_with_line_numbers, + "planning", on_chunk, on_complete ) diff --git a/lua/avante/utils/init.lua b/lua/avante/utils/init.lua index 33ad7ba..40fb6df 100644 --- a/lua/avante/utils/init.lua +++ b/lua/avante/utils/init.lua @@ -410,4 +410,10 @@ function M.is_type(type_name, v) end -- luacheck: pop +---@param code string +---@return string +function M.get_indentation(code) + return code:match("^%s*") or "" +end + return M