feat(tokens): add token count display to sidebar (#956)
* feat (tokens) add token count display to sidebar * refactor: calculate the real tokens and reuse input hints to avoid occlusion --------- Co-authored-by: yetone <yetoneful@gmail.com>
This commit is contained in:
Michael Gendy
@@ -1491,6 +1491,74 @@ function Sidebar:create_input_container(opts)
|
||||
|
||||
local chat_history = Path.history.load(self.code.bufnr)
|
||||
|
||||
---@param request string
|
||||
---@return GeneratePromptsOptions
|
||||
local function get_generate_prompts_options(request)
|
||||
local filetype = api.nvim_get_option_value("filetype", { buf = self.code.bufnr })
|
||||
|
||||
local selected_code_content = nil
|
||||
if self.code.selection ~= nil then selected_code_content = self.code.selection.content end
|
||||
|
||||
local mentions = Utils.extract_mentions(request)
|
||||
request = mentions.new_content
|
||||
|
||||
local file_ext = api.nvim_buf_get_name(self.code.bufnr):match("^.+%.(.+)$")
|
||||
|
||||
local project_context = mentions.enable_project_context and RepoMap.get_repo_map(file_ext) or nil
|
||||
|
||||
local selected_files_contents = self.file_selector:get_selected_files_contents()
|
||||
|
||||
local diagnostics = nil
|
||||
if mentions.enable_diagnostics then
|
||||
if self.code ~= nil and self.code.bufnr ~= nil and self.code.selection ~= nil then
|
||||
diagnostics = Utils.get_current_selection_diagnostics(self.code.bufnr, self.code.selection)
|
||||
else
|
||||
diagnostics = Utils.get_diagnostics(self.code.bufnr)
|
||||
end
|
||||
end
|
||||
|
||||
local history_messages = {}
|
||||
for i = #chat_history, 1, -1 do
|
||||
local entry = chat_history[i]
|
||||
if entry.reset_memory then break end
|
||||
if
|
||||
entry.request == nil
|
||||
or entry.original_response == nil
|
||||
or entry.request == ""
|
||||
or entry.original_response == ""
|
||||
then
|
||||
break
|
||||
end
|
||||
table.insert(history_messages, 1, { role = "assistant", content = entry.original_response })
|
||||
local user_content = ""
|
||||
if entry.selected_file ~= nil then
|
||||
user_content = user_content .. "SELECTED FILE: " .. entry.selected_file.filepath .. "\n\n"
|
||||
end
|
||||
if entry.selected_code ~= nil then
|
||||
user_content = user_content
|
||||
.. "SELECTED CODE:\n\n```"
|
||||
.. entry.selected_code.filetype
|
||||
.. "\n"
|
||||
.. entry.selected_code.content
|
||||
.. "\n```\n\n"
|
||||
end
|
||||
user_content = user_content .. "USER PROMPT:\n\n" .. entry.request
|
||||
table.insert(history_messages, 1, { role = "user", content = user_content })
|
||||
end
|
||||
|
||||
return {
|
||||
ask = opts.ask,
|
||||
project_context = vim.json.encode(project_context),
|
||||
selected_files = selected_files_contents,
|
||||
diagnostics = vim.json.encode(diagnostics),
|
||||
history_messages = history_messages,
|
||||
code_lang = filetype,
|
||||
selected_code = selected_code_content,
|
||||
instructions = request,
|
||||
mode = "planning",
|
||||
}
|
||||
end
|
||||
|
||||
---@param request string
|
||||
local function handle_submit(request)
|
||||
local model = Config.has_provider(Config.provider) and Config.get_provider(Config.provider).model or "default"
|
||||
@@ -1518,9 +1586,6 @@ function Sidebar:create_input_container(opts)
|
||||
self:update_content("", { focus = true, scroll = false })
|
||||
self:update_content(content_prefix .. generating_text)
|
||||
|
||||
local selected_code_content = nil
|
||||
if self.code.selection ~= nil then selected_code_content = self.code.selection.content end
|
||||
|
||||
if request:sub(1, 1) == "/" then
|
||||
local command, args = request:match("^/(%S+)%s*(.*)")
|
||||
if command == nil then
|
||||
@@ -1542,8 +1607,6 @@ function Sidebar:create_input_container(opts)
|
||||
Utils.error("Invalid end line number", { once = true, title = "Avante" })
|
||||
return
|
||||
end
|
||||
selected_code_content =
|
||||
table.concat(api.nvim_buf_get_lines(self.code.bufnr, start_line - 1, end_line, false), "\n")
|
||||
request = question
|
||||
end)
|
||||
else
|
||||
@@ -1632,67 +1695,15 @@ function Sidebar:create_input_container(opts)
|
||||
Path.history.save(self.code.bufnr, chat_history)
|
||||
end
|
||||
|
||||
local mentions = Utils.extract_mentions(request)
|
||||
request = mentions.new_content
|
||||
|
||||
local file_ext = api.nvim_buf_get_name(self.code.bufnr):match("^.+%.(.+)$")
|
||||
|
||||
local project_context = mentions.enable_project_context and RepoMap.get_repo_map(file_ext) or nil
|
||||
|
||||
local selected_files_contents = self.file_selector:get_selected_files_contents()
|
||||
|
||||
local diagnostics = nil
|
||||
if mentions.enable_diagnostics then
|
||||
if self.code ~= nil and self.code.bufnr ~= nil and self.code.selection ~= nil then
|
||||
diagnostics = Utils.get_current_selection_diagnostics(self.code.bufnr, self.code.selection)
|
||||
else
|
||||
diagnostics = Utils.get_diagnostics(self.code.bufnr)
|
||||
end
|
||||
end
|
||||
|
||||
local history_messages = {}
|
||||
for i = #chat_history, 1, -1 do
|
||||
local entry = chat_history[i]
|
||||
if entry.reset_memory then break end
|
||||
if
|
||||
entry.request == nil
|
||||
or entry.original_response == nil
|
||||
or entry.request == ""
|
||||
or entry.original_response == ""
|
||||
then
|
||||
break
|
||||
end
|
||||
table.insert(history_messages, 1, { role = "assistant", content = entry.original_response })
|
||||
local user_content = ""
|
||||
if entry.selected_file ~= nil then
|
||||
user_content = user_content .. "SELECTED FILE: " .. entry.selected_file.filepath .. "\n\n"
|
||||
end
|
||||
if entry.selected_code ~= nil then
|
||||
user_content = user_content
|
||||
.. "SELECTED CODE:\n\n```"
|
||||
.. entry.selected_code.filetype
|
||||
.. "\n"
|
||||
.. entry.selected_code.content
|
||||
.. "\n```\n\n"
|
||||
end
|
||||
user_content = user_content .. "USER PROMPT:\n\n" .. entry.request
|
||||
table.insert(history_messages, 1, { role = "user", content = user_content })
|
||||
end
|
||||
|
||||
Llm.stream({
|
||||
bufnr = self.code.bufnr,
|
||||
ask = opts.ask,
|
||||
project_context = vim.json.encode(project_context),
|
||||
selected_files = selected_files_contents,
|
||||
diagnostics = vim.json.encode(diagnostics),
|
||||
history_messages = history_messages,
|
||||
code_lang = filetype,
|
||||
selected_code = selected_code_content,
|
||||
instructions = request,
|
||||
mode = "planning",
|
||||
local generate_prompts_options = get_generate_prompts_options(request)
|
||||
---@type StreamOptions
|
||||
---@diagnostic disable-next-line: assign-type-mismatch
|
||||
local stream_options = vim.tbl_deep_extend("force", generate_prompts_options, {
|
||||
on_chunk = on_chunk,
|
||||
on_complete = on_complete,
|
||||
})
|
||||
|
||||
Llm.stream(stream_options)
|
||||
end
|
||||
|
||||
local get_position = function()
|
||||
@@ -1827,7 +1838,15 @@ function Sidebar:create_input_container(opts)
|
||||
local function show_hint()
|
||||
close_hint() -- Close the existing hint window
|
||||
|
||||
local hint_text = (fn.mode() ~= "i" and Config.mappings.submit.normal or Config.mappings.submit.insert)
|
||||
local input_value = table.concat(api.nvim_buf_get_lines(self.input_container.bufnr, 0, -1, false), "\n")
|
||||
|
||||
local generate_prompts_options = get_generate_prompts_options(input_value)
|
||||
local tokens = Llm.calculate_tokens(generate_prompts_options)
|
||||
|
||||
local hint_text = "Tokens: "
|
||||
.. tostring(tokens)
|
||||
.. "; "
|
||||
.. (fn.mode() ~= "i" and Config.mappings.submit.normal or Config.mappings.submit.insert)
|
||||
.. ": submit"
|
||||
|
||||
local buf = api.nvim_create_buf(false, true)
|
||||
|
||||
Reference in New Issue
Block a user