feat: automatic suggestion (smart tab) (yetone#455)

ogfunkycold · Sep 3, 2024 · 65e1e17 · 65e1e17
1 parent 962dd0a
commit 65e1e17
Show file tree

Hide file tree

Showing 10 changed files with 592 additions and 43 deletions.
diff --git a/README.md b/README.md
@@ -192,6 +192,13 @@ _See [config.lua#L9](./lua/avante/config.lua) for the full config_
     temperature = 0,
     max_tokens = 4096,
   },
+  behaviour = {
+    auto_suggestions = false, -- Experimental stage
+    auto_set_highlight_group = true,
+    auto_set_keymaps = true,
+    auto_apply_diff_after_generation = false,
+    support_paste_from_clipboard = false,
+  },
   mappings = {
     --- @class AvanteConflictMappings
     diff = {
@@ -203,6 +210,12 @@ _See [config.lua#L9](./lua/avante/config.lua) for the full config_
       next = "]x",
       prev = "[x",
     },
+    suggestion = {
+      accept = "<M-l>",
+      next = "<M-]>",
+      prev = "<M-[>",
+      dismiss = "<C-]>",
+    },
     jump = {
       next = "]]",
       prev = "[[",

diff --git a/lua/avante/api.lua b/lua/avante/api.lua
@@ -45,8 +45,14 @@ M.edit = function(question)
   end
 end
 
+---@return avante.Suggestion | nil
+M.get_suggestion = function()
+  local _, _, suggestion = require("avante").get()
+  return suggestion
+end
+
 M.refresh = function()
-  local sidebar, _ = require("avante").get()
+  local sidebar = require("avante").get()
   if not sidebar then
     return
   end

diff --git a/lua/avante/config.lua b/lua/avante/config.lua
@@ -86,6 +86,7 @@ M.defaults = {
   ---3. auto_set_highlight_group        : Whether to automatically set the highlight group for the current line. Default to true.
   ---4. support_paste_from_clipboard    : Whether to support pasting image from clipboard. This will be determined automatically based whether img-clip is available or not.
   behaviour = {
+    auto_suggestions = false, -- Experimental stage
     auto_set_highlight_group = true,
     auto_set_keymaps = true,
     auto_apply_diff_after_generation = false,
@@ -116,6 +117,12 @@ M.defaults = {
       next = "]x",
       prev = "[x",
     },
+    suggestion = {
+      accept = "<M-l>",
+      next = "<M-]>",
+      prev = "<M-[>",
+      dismiss = "<C-]>",
+    },
     jump = {
       next = "]]",
       prev = "[[",

diff --git a/lua/avante/highlights.lua b/lua/avante/highlights.lua
@@ -11,6 +11,8 @@ local Highlights = {
   REVERSED_SUBTITLE = { name = "AvanteReversedSubtitle", fg = "#56b6c2" },
   THIRD_TITLE = { name = "AvanteThirdTitle", fg = "#ABB2BF", bg = "#353B45" },
   REVERSED_THIRD_TITLE = { name = "AvanteReversedThirdTitle", fg = "#353B45" },
+  SUGGESTION = { name = "AvanteSuggestion", link = "Comment" },
+  ANNOTATION = { name = "AvanteAnnotation", link = "Comment" },
 }
 
 Highlights.conflict = {
@@ -48,7 +50,7 @@ M.setup = function()
       end)
       :each(function(_, hl)
         if not has_set_colors(hl.name) then
-          api.nvim_set_hl(0, hl.name, { fg = hl.fg or nil, bg = hl.bg or nil })
+          api.nvim_set_hl(0, hl.name, { fg = hl.fg or nil, bg = hl.bg or nil, link = hl.link or nil })
         end
       end)
 

diff --git a/lua/avante/init.lua b/lua/avante/init.lua
@@ -3,6 +3,7 @@ local api = vim.api
 local Utils = require("avante.utils")
 local Sidebar = require("avante.sidebar")
 local Selection = require("avante.selection")
+local Suggestion = require("avante.suggestion")
 local Config = require("avante.config")
 local Diff = require("avante.diff")
 
@@ -12,8 +13,10 @@ local M = {
   sidebars = {},
   ---@type avante.Selection[]
   selections = {},
-  ---@type {sidebar?: avante.Sidebar, selection?: avante.Selection}
-  current = { sidebar = nil, selection = nil },
+  ---@type avante.Suggestion[]
+  suggestions = {},
+  ---@type {sidebar?: avante.Sidebar, selection?: avante.Selection, suggestion?: avante.Suggestion}
+  current = { sidebar = nil, selection = nil, suggestion = nil },
 }
 
 M.did_setup = false
@@ -185,7 +188,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
@@ -234,22 +237,25 @@ H.autocmds = function()
 end
 
 ---@param current boolean? false to disable setting current, otherwise use this to track across tabs.
----@return avante.Sidebar, avante.Selection
+---@return avante.Sidebar, avante.Selection, avante.Suggestion
 function M.get(current)
   local tab = api.nvim_get_current_tabpage()
   local sidebar = M.sidebars[tab]
   local selection = M.selections[tab]
+  local suggestion = M.suggestions[tab]
   if current ~= false then
     M.current.sidebar = sidebar
     M.current.selection = selection
+    M.current.suggestion = suggestion
   end
-  return sidebar, selection
+  return sidebar, selection, suggestion
 end
 
 ---@param id integer
 function M._init(id)
   local sidebar = M.sidebars[id]
   local selection = M.selections[id]
+  local suggestion = M.suggestions[id]
 
   if not sidebar then
     sidebar = Sidebar:new(id)
@@ -259,7 +265,11 @@ function M._init(id)
     selection = Selection:new(id)
     M.selections[id] = selection
   end
-  M.current = { sidebar = sidebar, selection = selection }
+  if not suggestion then
+    suggestion = Suggestion:new(id)
+    M.suggestions[id] = suggestion
+  end
+  M.current = { sidebar = sidebar, selection = selection, suggestion = suggestion }
   return M
 end
 
@@ -288,7 +298,7 @@ M.toggle.hint = H.api(Utils.toggle_wrap({
 setmetatable(M.toggle, {
   __index = M.toggle,
   __call = function()
-    local sidebar, _ = M.get()
+    local sidebar = M.get()
     if not sidebar then
       M._init(api.nvim_get_current_tabpage())
       M.current.sidebar:open()
@@ -326,15 +336,15 @@ M.build = H.api(function()
   local os_name = Utils.get_os_name()
 
   if vim.tbl_contains({ "linux", "darwin" }, os_name) then
-    cmd = { "sh", "-c", ("make -C %s"):format(build_directory) }
+    cmd = { "sh", "-c", string.format("make -C %s", build_directory) }
   elseif os_name == "windows" then
     build_directory = to_windows_path(build_directory)
     cmd = {
       "powershell",
       "-ExecutionPolicy",
       "Bypass",
       "-File",
-      ("%s\\Build.ps1"):format(build_directory),
+      string.format("%s\\Build.ps1", build_directory),
       "-WorkingDirectory",
       build_directory,
     }

diff --git a/lua/avante/llm.lua b/lua/avante/llm.lua
@@ -88,8 +88,36 @@ Your task is to modify the provided code according to the user's request. Follow
 Remember: Your response should contain ONLY the modified code, ready to be used as a direct replacement for the original file.
 ]]
 
+local suggesting_mode_user_prompt_tpl = [[
+Your task is to suggest code modifications at the cursor position. Follow these instructions meticulously:
+
+1. Carefully analyze the original code, paying close attention to its structure and the cursor position.
+
+2. You must follow this json format when suggesting modifications:
+
+[
+  {
+    "row": ${row},
+    "col": ${column},
+    "content": "Your suggested code here"
+  }
+]
+
+3. When suggesting suggested code:
+  - Each element in the returned list is a COMPLETE and INDEPENDENT code snippet.
+  - MUST be a valid json format. Don't be lazy!
+  - Only return the new code to be inserted.
+  - Your returned code should not overlap with the original code in any way. Don't be lazy!
+  - Please strictly check the code around the position and ensure that the complete code after insertion is correct. Don't be lazy!
+  - Do not return the entire file content or any surrounding code.
+  - Do not include any explanations, comments, or line numbers in your response.
+  - Ensure the suggested code fits seamlessly with the existing code structure and indentation.
+  - If there are no recommended modifications, return an empty list.
+
+Remember: Return ONLY the suggested code snippet, without any additional formatting or explanation.
+]]
+
 local group = api.nvim_create_augroup("avante_llm", { clear = true })
-local active_job = nil
 
 ---@class StreamOptions
 ---@field file_content string
@@ -99,7 +127,7 @@ local active_job = nil
 ---@field project_context string | nil
 ---@field memory_context string | nil
 ---@field full_file_contents_context string | nil
----@field mode "planning" | "editing"
+---@field mode "planning" | "editing" | "suggesting"
 ---@field on_chunk AvanteChunkParser
 ---@field on_complete AvanteCompleteParser
 
@@ -108,7 +136,13 @@ M.stream = function(opts)
   local mode = opts.mode or "planning"
   local provider = Config.provider
 
-  local user_prompt_tpl = mode == "planning" and planning_mode_user_prompt_tpl or editing_mode_user_prompt_tpl
+  local user_prompt_tpl = planning_mode_user_prompt_tpl
+
+  if mode == "editing" then
+    user_prompt_tpl = editing_mode_user_prompt_tpl
+  elseif mode == "suggesting" then
+    user_prompt_tpl = suggesting_mode_user_prompt_tpl
+  end
 
   -- Check if the instructions contains an image path
   local image_paths = {}
@@ -191,13 +225,10 @@ M.stream = function(opts)
     end
   end
 
-  if active_job then
-    active_job:shutdown()
-    active_job = nil
-  end
-
   local completed = false
 
+  local active_job
+
   active_job = curl.post(spec.url, {
     headers = spec.headers,
     proxy = spec.proxy,
@@ -230,11 +261,13 @@ M.stream = function(opts)
         end
       end)
     end,
-    on_error = function(err)
+    on_error = function()
+      active_job = nil
       completed = true
-      opts.on_complete(err)
+      opts.on_complete(nil)
     end,
     callback = function(result)
+      active_job = nil
       if result.status >= 400 then
         if Provider.on_error then
           Provider.on_error(result)
@@ -250,16 +283,21 @@ M.stream = function(opts)
           end
         end)
       end
-      active_job = nil
     end,
   })
 
   api.nvim_create_autocmd("User", {
     group = group,
     pattern = M.CANCEL_PATTERN,
+    once = true,
     callback = function()
+      -- Error: cannot resume dead coroutine
       if active_job then
-        active_job:shutdown()
+        xpcall(function()
+          active_job:shutdown()
+        end, function(err)
+          return err
+        end)
         Utils.debug("LLM request cancelled", { title = "Avante" })
         active_job = nil
       end
@@ -269,4 +307,8 @@ M.stream = function(opts)
   return active_job
 end
 
+function M.cancel_inflight_request()
+  api.nvim_exec_autocmds("User", { pattern = M.CANCEL_PATTERN })
+end
+
 return M
diff --git a/lua/avante/selection.lua b/lua/avante/selection.lua
@@ -84,7 +84,7 @@ end
 
 function Selection:close_editing_input()
   self:close_editing_input_shortcuts_hints()
-  api.nvim_exec_autocmds("User", { pattern = Llm.CANCEL_PATTERN })
+  Llm.cancel_inflight_request()
   if api.nvim_get_mode().mode == "i" then
     vim.cmd([[stopinsert]])
   end

diff --git a/lua/avante/sidebar.lua b/lua/avante/sidebar.lua
@@ -448,7 +448,7 @@ local function insert_conflict_contents(bufnr, snippets)
     local snippet_lines = vim.split(snippet.content, "\n")
 
     for idx, line in ipairs(snippet_lines) do
-      line = line:gsub("^L%d+: ", "")
+      line = Utils.trim_line_number(line)
       if idx == 1 then
         local indentation = Utils.get_indentation(line)
         need_prepend_indentation = indentation ~= original_start_line_indentation
@@ -824,7 +824,7 @@ function Sidebar:on_mount()
   self:render_input()
   self:render_selected_code()
 
-  self.augroup = api.nvim_create_augroup("avante_" .. self.id .. self.result.winid, { clear = true })
+  self.augroup = api.nvim_create_augroup("avante_sidebar_" .. self.id .. self.result.winid, { clear = true })
 
   local filetype = api.nvim_get_option_value("filetype", { buf = self.code.bufnr })
 
@@ -1040,17 +1040,6 @@ function Sidebar:update_content(content, opts)
   return self
 end
 
-local function prepend_line_number(content, start_line)
-  start_line = start_line or 1
-  local lines = vim.split(content, "\n")
-  local result = {}
-  for i, line in ipairs(lines) do
-    i = i + start_line - 1
-    table.insert(result, "L" .. i .. ": " .. line)
-  end
-  return table.concat(result, "\n")
-end
-
 -- Function to get current timestamp
 local function get_timestamp()
   return os.date("%Y-%m-%d %H:%M:%S")
@@ -1253,14 +1242,14 @@ function Sidebar:create_input()
     self:update_content(content_prefix .. "🔄 **Generating response ...**\n")
 
     local content = table.concat(Utils.get_buf_lines(0, -1, self.code.bufnr), "\n")
-    local content_with_line_numbers = prepend_line_number(content)
+    local content_with_line_numbers = Utils.prepend_line_number(content)
 
     local filetype = api.nvim_get_option_value("filetype", { buf = self.code.bufnr })
 
     local selected_code_content_with_line_numbers = nil
     if self.code.selection ~= nil then
       selected_code_content_with_line_numbers =
-        prepend_line_number(self.code.selection.content, self.code.selection.range.start.line)
+        Utils.prepend_line_number(self.code.selection.content, self.code.selection.range.start.line)
     end
 
     if request:sub(1, 1) == "/" then
@@ -1289,7 +1278,7 @@ function Sidebar:create_input()
               Utils.error("Invalid end line number", { once = true, title = "Avante" })
               return
             end
-            selected_code_content_with_line_numbers = prepend_line_number(
+            selected_code_content_with_line_numbers = Utils.prepend_line_number(
               table.concat(api.nvim_buf_get_lines(self.code.bufnr, start_line - 1, end_line, false), "\n"),
               start_line
             )
@@ -1640,12 +1629,12 @@ function Sidebar:render()
   end)
 
   self.result:map("n", "q", function()
-    api.nvim_exec_autocmds("User", { pattern = Llm.CANCEL_PATTERN })
+    Llm.cancel_inflight_request()
     self:close()
   end)
 
   self.result:map("n", "<Esc>", function()
-    api.nvim_exec_autocmds("User", { pattern = Llm.CANCEL_PATTERN })
+    Llm.cancel_inflight_request()
     self:close()
   end)