From 2ead26f809dd9804678de3dd18fa65ab3fadce29 Mon Sep 17 00:00:00 2001 From: yetone Date: Mon, 27 Jan 2025 22:50:40 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20thinking=20=F0=9F=A4=AF?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- lua/avante/sidebar.lua | 64 ++++++++++++++++++++++++++++++++++++++---- 1 file changed, 59 insertions(+), 5 deletions(-) diff --git a/lua/avante/sidebar.lua b/lua/avante/sidebar.lua index 5a9f21f..535a882 100644 --- a/lua/avante/sidebar.lua +++ b/lua/avante/sidebar.lua @@ -170,11 +170,15 @@ end ---@field current_filepath string ---@field is_searching boolean ---@field is_replacing boolean +---@field is_thinking boolean ---@field last_search_tag_start_line integer ---@field last_replace_tag_start_line integer +---@field last_think_tag_start_line integer +---@field last_think_tag_end_line integer ---@param selected_files {path: string, content: string, file_type: string | nil}[] ---@param result_content string +---@param prev_filepath string ---@return AvanteReplacementResult local function transform_result_content(selected_files, result_content, prev_filepath) local transformed_lines = {} @@ -183,8 +187,11 @@ local function transform_result_content(selected_files, result_content, prev_fil local is_searching = false local is_replacing = false + local is_thinking = false local last_search_tag_start_line = 0 local last_replace_tag_start_line = 0 + local last_think_tag_start_line = 0 + local last_think_tag_end_line = 0 local search_start = 0 @@ -286,6 +293,12 @@ local function transform_result_content(selected_files, result_content, prev_fil local prev_line = result_lines[i - 1] if not (prev_line and prev_line:match("^%s*```$")) then table.insert(transformed_lines, "```") end goto continue + elseif line_content == "" then + is_thinking = true + last_think_tag_start_line = i + elseif line_content == "" then + is_thinking = false + last_think_tag_end_line = i end table.insert(transformed_lines, line_content) ::continue:: @@ -297,8 +310,11 @@ local function transform_result_content(selected_files, result_content, prev_fil content = table.concat(transformed_lines, "\n"), is_searching = is_searching, is_replacing = is_replacing, + is_thinking = is_thinking, last_search_tag_start_line = last_search_tag_start_line, last_replace_tag_start_line = last_replace_tag_start_line, + last_think_tag_start_line = last_think_tag_start_line, + last_think_tag_end_line = last_think_tag_end_line, } end @@ -351,8 +367,22 @@ local function get_searching_hint() return "\n" .. spinner .. " Searching..." end +local thinking_spinner_chars = { + "🤯", + "🙄", +} +local thinking_spinner_index = 1 + +local function get_thinking_spinner() + thinking_spinner_index = thinking_spinner_index + 1 + if thinking_spinner_index > #thinking_spinner_chars then thinking_spinner_index = 1 end + local spinner = thinking_spinner_chars[thinking_spinner_index] + return "\n\n" .. spinner .. " Thinking..." +end + local function get_display_content_suffix(replacement) if replacement.is_searching then return get_searching_hint() end + if replacement.is_thinking then return get_thinking_spinner() end return "" end @@ -365,6 +395,25 @@ local function generate_display_content(replacement) "\n" ) end + if replacement.last_think_tag_start_line > 0 then + local lines = vim.split(replacement.content, "\n") + local last_think_tag_end_line = replacement.last_think_tag_end_line + if last_think_tag_end_line == 0 then last_think_tag_end_line = #lines + 1 end + local thinking_content_lines = + vim.list_slice(lines, replacement.last_think_tag_start_line + 2, last_think_tag_end_line - 1) + local formatted_thinking_content_lines = vim + .iter(thinking_content_lines) + :map(function(line) + if Utils.trim_spaces(line) == "" then return line end + return string.format(" > %s", line) + end) + :totable() + local result_lines = + vim.list_extend(vim.list_slice(lines, 1, replacement.last_search_tag_start_line), { "🤔 Thought content:" }) + result_lines = vim.list_extend(result_lines, formatted_thinking_content_lines) + result_lines = vim.list_extend(result_lines, vim.list_slice(lines, last_think_tag_end_line + 1)) + return table.concat(result_lines, "\n") + end return replacement.content end @@ -1675,6 +1724,7 @@ function Sidebar:create_input_container(opts) local transformed_response = "" local displayed_response = "" local current_path = "" + local prev_is_thinking = false local is_first_chunk = true local scroll = true @@ -1708,8 +1758,10 @@ function Sidebar:create_input_container(opts) local selected_files = self.file_selector:get_selected_files_contents() - local transformed = transform_result_content(selected_files, transformed_response .. chunk, current_path) + local transformed = + transform_result_content(selected_files, transformed_response .. chunk, current_path, prev_is_thinking) transformed_response = transformed.content + prev_is_thinking = transformed.is_thinking if transformed.current_filepath and transformed.current_filepath ~= "" then current_path = transformed.current_filepath end @@ -1727,10 +1779,12 @@ function Sidebar:create_input_container(opts) ---@type AvanteCompleteParser local on_complete = function(err) - ---remove keymaps - vim.keymap.del("n", "j", { buffer = self.result_container.bufnr }) - vim.keymap.del("n", "k", { buffer = self.result_container.bufnr }) - vim.keymap.del("n", "G", { buffer = self.result_container.bufnr }) + pcall(function() + ---remove keymaps + vim.keymap.del("n", "j", { buffer = self.result_container.bufnr }) + vim.keymap.del("n", "k", { buffer = self.result_container.bufnr }) + vim.keymap.del("n", "G", { buffer = self.result_container.bufnr }) + end) if err ~= nil then self:update_content(