feat: custom slash commands (#1826)

lua/avante/config.lua

@@ -492,6 +492,8 @@ M._defaults = {
   disabled_tools = {}, ---@type string[]
   ---@type AvanteLLMToolPublic[] | fun(): AvanteLLMToolPublic[]
   custom_tools = {},
+  ---@type AvanteSlashCommand[]
+  slash_commands = {},
 }
 
 ---@type avante.Config

lua/avante/init.lua

@@ -447,6 +447,59 @@ function M.setup(opts)
   end
 
   if Config.rag_service.enabled then run_rag_service() end
+
+  local has_cmp, cmp = pcall(require, "cmp")
+  if has_cmp then
+    cmp.register_source("avante_commands", require("cmp_avante.commands"):new())
+
+    cmp.register_source(
+      "avante_mentions",
+      require("cmp_avante.mentions"):new(function()
+        local mentions = Utils.get_mentions()
+
+        table.insert(mentions, {
+          description = "file",
+          command = "file",
+          details = "add files...",
+          callback = function(sidebar) sidebar.file_selector:open() end,
+        })
+
+        table.insert(mentions, {
+          description = "quickfix",
+          command = "quickfix",
+          details = "add files in quickfix list to chat context",
+          callback = function(sidebar) sidebar.file_selector:add_quickfix_files() end,
+        })
+
+        table.insert(mentions, {
+          description = "buffers",
+          command = "buffers",
+          details = "add open buffers to the chat context",
+          callback = function(sidebar) sidebar.file_selector:add_buffer_files() end,
+        })
+
+        return mentions
+      end)
+    )
+
+    cmp.register_source("avante_prompt_mentions", require("cmp_avante.mentions"):new(Utils.get_mentions))
+
+    cmp.setup.filetype({ "AvanteInput" }, {
+      enabled = true,
+      sources = {
+        { name = "avante_commands" },
+        { name = "avante_mentions" },
+        { name = "avante_files" },
+      },
+    })
+
+    cmp.setup.filetype({ "AvantePromptInput" }, {
+      enabled = true,
+      sources = {
+        { name = "avante_prompt_mentions" },
+      },
+    })
+  end
 end
 
 return M

lua/avante/llm.lua

@@ -155,13 +155,6 @@ end
 ---@param opts AvanteGeneratePromptsOptions
 ---@return AvantePromptOptions
 function M.generate_prompts(opts)
-  if opts.prompt_opts then
-    local prompt_opts = vim.tbl_deep_extend("force", opts.prompt_opts, {
-      tool_histories = opts.tool_histories,
-    })
-    ---@cast prompt_opts AvantePromptOptions
-    return prompt_opts
-  end
   local provider = opts.provider or Providers[Config.provider]
   local mode = opts.mode or "planning"
   ---@type AvanteProviderFunctor | AvanteBedrockProviderFunctor
@@ -170,6 +163,9 @@ function M.generate_prompts(opts)
 
   -- Check if the instructions contains an image path
   local image_paths = {}
+  if opts.prompt_opts and opts.prompt_opts.image_paths then
+    image_paths = vim.list_extend(image_paths, opts.prompt_opts.image_paths)
+  end
   local instructions = opts.instructions
   if instructions and instructions:match("image: ") then
     local lines = vim.split(opts.instructions, "\n")
@@ -201,7 +197,12 @@ function M.generate_prompts(opts)
     memory = opts.memory,
   }
 
-  local system_prompt = Path.prompts.render_mode(mode, template_opts)
+  local system_prompt
+  if opts.prompt_opts and opts.prompt_opts.system_prompt then
+    system_prompt = opts.prompt_opts.system_prompt
+  else
+    system_prompt = Path.prompts.render_mode(mode, template_opts)
+  end
 
   if Config.system_prompt ~= nil then
     local custom_system_prompt = Config.system_prompt
@@ -213,6 +214,9 @@ function M.generate_prompts(opts)
 
   ---@type AvanteLLMMessage[]
   local messages = {}
+  if opts.prompt_opts and opts.prompt_opts.messages then
+    messages = vim.list_extend(messages, opts.prompt_opts.messages)
+  end
 
   if opts.project_context ~= nil and opts.project_context ~= "" and opts.project_context ~= "null" then
     local project_context = Path.prompts.render_file("_project.avanterules", template_opts)
@@ -243,6 +247,10 @@ function M.generate_prompts(opts)
   end
 
   local dropped_history_messages = {}
+  if opts.prompt_opts and opts.prompt_opts.dropped_history_messages then
+    dropped_history_messages = vim.list_extend(dropped_history_messages, opts.prompt_opts.dropped_history_messages)
+  end
+
   if opts.history_messages then
     if Config.history.max_tokens > 0 then remaining_tokens = math.min(Config.history.max_tokens, remaining_tokens) end
     -- Traverse the history in reverse, keeping only the latest history until the remaining tokens are exhausted and the first message role is "user"
@@ -290,13 +298,23 @@ Merge all changes from the <update> snippet into the <code> below.
   opts.session_ctx.system_prompt = system_prompt
   opts.session_ctx.messages = messages
 
+  local tools = {}
+  if opts.tools then tools = vim.list_extend(tools, opts.tools) end
+  if opts.prompt_opts and opts.prompt_opts.tools then tools = vim.list_extend(tools, opts.prompt_opts.tools) end
+
+  local tool_histories = {}
+  if opts.tool_histories then tool_histories = vim.list_extend(tool_histories, opts.tool_histories) end
+  if opts.prompt_opts and opts.prompt_opts.tool_histories then
+    tool_histories = vim.list_extend(tool_histories, opts.prompt_opts.tool_histories)
+  end
+
   ---@type AvantePromptOptions
   return {
     system_prompt = system_prompt,
     messages = messages,
     image_paths = image_paths,
-    tools = opts.tools,
+    tools = tools,
-    tool_histories = opts.tool_histories,
+    tool_histories = tool_histories,
     dropped_history_messages = dropped_history_messages,
   }
 end

lua/avante/selection.lua

@@ -292,28 +292,6 @@ function Selection:create_editing_input(request, line1, line2)
   self.prompt_input = prompt_input
 
   prompt_input:open()
-
-  api.nvim_create_autocmd("InsertEnter", {
-    group = self.augroup,
-    buffer = prompt_input.bufnr,
-    once = true,
-    desc = "Setup the completion of helpers in the input buffer",
-    callback = function()
-      local has_cmp, cmp = pcall(require, "cmp")
-      if has_cmp then
-        cmp.register_source(
-          "avante_mentions",
-          require("cmp_avante.mentions"):new(Utils.get_mentions(), prompt_input.bufnr)
-        )
-        cmp.setup.buffer({
-          enabled = true,
-          sources = {
-            { name = "avante_mentions" },
-          },
-        })
-      end
-    end,
-  })
 end
 
 function Selection:setup_autocmds()

lua/avante/sidebar.lua

@@ -635,6 +635,10 @@ local function tree_sitter_markdown_parse_code_blocks(source)
   else
     parser = vim.treesitter.get_parser(source, "markdown")
   end
+  if parser == nil then
+    Utils.warn("Failed to get markdown parser")
+    return {}
+  end
   local tree = parser:parse()[1]
   local root = tree:root()
   local code_block_query = query.parse(
@@ -2203,29 +2207,46 @@ end
 ---@param history avante.ChatHistory
 ---@return string
 function Sidebar.render_history_content(history)
+  local added_breakline = false
   local content = ""
   for idx, entry in ipairs(history.entries) do
     if entry.visible == false then goto continue end
     if entry.reset_memory then
       content = content .. "***MEMORY RESET***\n\n"
-      if idx < #history.entries then content = content .. "-------\n\n" end
+      if idx < #history.entries and not added_breakline then
+        added_breakline = true
+        content = content .. "-------\n\n"
+      end
       goto continue
     end
     local selected_filepaths = entry.selected_filepaths
     if not selected_filepaths and entry.selected_file ~= nil then
       selected_filepaths = { entry.selected_file.filepath }
     end
-    local prefix = render_chat_record_prefix(
+    if entry.request and entry.request ~= "" then
-      entry.timestamp,
+      if idx ~= 1 and not added_breakline then
-      entry.provider,
+        added_breakline = true
-      entry.model,
+        content = content .. "-------\n\n"
-      entry.request or "",
+      end
-      selected_filepaths or {},
+      local prefix = render_chat_record_prefix(
-      entry.selected_code
+        entry.timestamp,
-    )
+        entry.provider,
-    content = content .. prefix
+        entry.model,
-    content = content .. entry.response .. "\n\n"
+        entry.request or "",
-    if idx < #history.entries then content = content .. "-------\n\n" end
+        selected_filepaths or {},
+        entry.selected_code
+      )
+      content = content .. prefix
+    end
+    if entry.response and entry.response ~= "" then
+      content = content .. entry.response .. "\n\n"
+      if idx < #history.entries then
+        added_breakline = true
+        content = content .. "-------\n\n"
+      end
+    else
+      added_breakline = false
+    end
     ::continue::
   end
   return content
@@ -2301,26 +2322,38 @@ function Sidebar:new_chat(args, cb)
   if cb then cb(args) end
 end
 
----@param message AvanteLLMMessage
+---@param messages AvanteLLMMessage | AvanteLLMMessage[]
 ---@param options {visible?: boolean}
-function Sidebar:add_chat_history(message, options)
+function Sidebar:add_chat_history(messages, options)
+  options = options or {}
   local timestamp = get_timestamp()
+  messages = vim.islist(messages) and messages or { messages }
   self:reload_chat_history()
-  table.insert(self.chat_history.entries, {
+  for _, message in ipairs(messages) do
-    timestamp = timestamp,
+    local content = message.content
-    provider = Config.provider,
+    if message.role == "system" and type(content) == "string" then
-    model = Config.get_provider_config(Config.provider).model,
+      ---@cast content string
-    request = message.role == "user" and message.content or "",
+      self.chat_history.system_prompt = content
-    response = message.role == "assistant" and message.content or "",
+      goto continue
-    original_response = "",
+    end
-    selected_filepaths = nil,
+    table.insert(self.chat_history.entries, {
-    selected_code = nil,
+      timestamp = timestamp,
-    reset_memory = false,
+      provider = Config.provider,
-    visible = options.visible,
+      model = Config.get_provider_config(Config.provider).model,
-  })
+      request = message.role == "user" and message.content or "",
+      response = message.role == "assistant" and message.content or "",
+      original_response = message.role == "assistant" and message.content or "",
+      selected_filepaths = nil,
+      selected_code = nil,
+      reset_memory = false,
+      visible = options.visible,
+    })
+    ::continue::
+  end
   Path.history.save(self.code.bufnr, self.chat_history)
-  if self.chat_history.title == "untitled" then
+  if options.visible then self:update_content_with_history() end
-    Llm.summarize_chat_thread_title(message.content, function(title)
+  if self.chat_history.title == "untitled" and #messages > 0 then
+    Llm.summarize_chat_thread_title(messages[1].content, function(title)
       self:reload_chat_history()
       if title then self.chat_history.title = title end
       Path.history.save(self.code.bufnr, self.chat_history)
@@ -2360,70 +2393,6 @@ function Sidebar:reset_memory(args, cb)
   end
 end
 
----@alias AvanteSlashCommandType "clear" | "help" | "lines" | "reset" | "commit" | "new"
----@alias AvanteSlashCommandCallback fun(args: string, cb?: fun(args: string): nil): nil
----@alias AvanteSlashCommand {description: string, command: AvanteSlashCommandType, details: string, shorthelp?: string, callback?: AvanteSlashCommandCallback}
----@return AvanteSlashCommand[]
-function Sidebar:get_commands()
-  ---@param items_ {command: string, description: string, shorthelp?: string}[]
-  ---@return string
-  local function get_help_text(items_)
-    local help_text = ""
-    for _, item in ipairs(items_) do
-      help_text = help_text .. "- " .. item.command .. ": " .. (item.shorthelp or item.description) .. "\n"
-    end
-    return help_text
-  end
-
-  ---@type AvanteSlashCommand[]
-  local items = {
-    { description = "Show help message", command = "help" },
-    { description = "Clear chat history", command = "clear" },
-    { description = "Reset memory", command = "reset" },
-    { description = "New chat", command = "new" },
-    {
-      shorthelp = "Ask a question about specific lines",
-      description = "/lines <start>-<end> <question>",
-      command = "lines",
-    },
-    { description = "Commit the changes", command = "commit" },
-  }
-
-  ---@type {[AvanteSlashCommandType]: AvanteSlashCommandCallback}
-  local cbs = {
-    help = function(args, cb)
-      local help_text = get_help_text(items)
-      self:update_content(help_text, { focus = false, scroll = false })
-      if cb then cb(args) end
-    end,
-    clear = function(args, cb) self:clear_history(args, cb) end,
-    reset = function(args, cb) self:reset_memory(args, cb) end,
-    new = function(args, cb) self:new_chat(args, cb) end,
-    lines = function(args, cb)
-      if cb then cb(args) end
-    end,
-    commit = function(_, cb)
-      local question = "Please commit the changes"
-      if cb then cb(question) end
-    end,
-  }
-
-  return vim
-    .iter(items)
-    :map(
-      ---@param item AvanteSlashCommand
-      function(item)
-        return {
-          command = item.command,
-          description = item.description,
-          callback = cbs[item.command],
-          details = item.shorthelp and table.concat({ item.shorthelp, item.description }, "\n") or item.description,
-        }
-      end
-    )
-    :totable()
-end
-
 function Sidebar:create_selected_code_container()
   if self.selected_code_container ~= nil then
     self.selected_code_container:unmount()
@@ -2576,6 +2545,13 @@ function Sidebar:create_input_container(opts)
       tools = tools,
     }
 
+    if self.chat_history.system_prompt then
+      prompts_opts.prompt_opts = {
+        system_prompt = self.chat_history.system_prompt,
+        messages = {},
+      }
+    end
+
     if self.chat_history.memory then prompts_opts.memory = self.chat_history.memory.content end
 
     if not summarize_memory or #history_messages < 8 then
@@ -2598,25 +2574,26 @@ function Sidebar:create_input_container(opts)
       return
     end
 
-    if request:sub(1, 1) == "/" then
+    local has_cmp = pcall(require, "cmp")
+    if request:sub(1, 1) == "/" and not has_cmp then
       local command, args = request:match("^/(%S+)%s*(.*)")
       if command == nil then
         self:update_content("Invalid command", { focus = false, scroll = false })
         return
       end
-      local cmds = self:get_commands()
+      local cmds = Utils.get_commands()
       ---@type AvanteSlashCommand
-      local cmd = vim.iter(cmds):filter(function(_) return _.command == command end):totable()[1]
+      local cmd = vim.iter(cmds):filter(function(cmd) return cmd.name == command end):totable()[1]
       if cmd then
         if command == "lines" then
-          cmd.callback(args, function(args_)
+          cmd.callback(self, args, function(args_)
             local _, _, question = args_:match("(%d+)-(%d+)%s+(.*)")
             request = question
           end)
         elseif command == "commit" then
-          cmd.callback(args, function(question) request = question end)
+          cmd.callback(self, args, function(question) request = question end)
         else
-          cmd.callback(args)
+          cmd.callback(self, args)
           return
         end
       else
@@ -2940,51 +2917,7 @@ function Sidebar:create_input_container(opts)
     buffer = self.input_container.bufnr,
     once = true,
     desc = "Setup the completion of helpers in the input buffer",
-    callback = function()
+    callback = function() end,
-      local has_cmp, cmp = pcall(require, "cmp")
-      if has_cmp then
-        local mentions = Utils.get_mentions()
-
-        table.insert(mentions, {
-          description = "file",
-          command = "file",
-          details = "add files...",
-          callback = function() self.file_selector:open() end,
-        })
-
-        table.insert(mentions, {
-          description = "quickfix",
-          command = "quickfix",
-          details = "add files in quickfix list to chat context",
-          callback = function() self.file_selector:add_quickfix_files() end,
-        })
-
-        table.insert(mentions, {
-          description = "buffers",
-          command = "buffers",
-          details = "add open buffers to the chat context",
-          callback = function() self.file_selector:add_buffer_files() end,
-        })
-
-        cmp.register_source(
-          "avante_commands",
-          require("cmp_avante.commands"):new(self:get_commands(), self.input_container.bufnr)
-        )
-        cmp.register_source(
-          "avante_mentions",
-          require("cmp_avante.mentions"):new(mentions, self.input_container.bufnr)
-        )
-
-        cmp.setup.buffer({
-          enabled = true,
-          sources = {
-            { name = "avante_commands" },
-            { name = "avante_mentions" },
-            { name = "avante_files" },
-          },
-        })
-      end
-    end,
   })
 
   -- Close the floating window
@@ -3130,6 +3063,20 @@ function Sidebar:create_input_container(opts)
   self:refresh_winids()
 end
 
+---@param value string
+function Sidebar:set_input_value(value)
+  if not self.input_container then return end
+  if not value then return end
+  api.nvim_buf_set_lines(self.input_container.bufnr, 0, -1, false, vim.split(value, "\n"))
+end
+
+---@return string
+function Sidebar:get_input_value()
+  if not self.input_container then return "" end
+  local lines = api.nvim_buf_get_lines(self.input_container.bufnr, 0, -1, false)
+  return table.concat(lines, "\n")
+end
+
 function Sidebar:get_selected_code_size()
   local selected_code_max_lines_count = 10
 

lua/avante/types.lua

@@ -398,6 +398,7 @@ vim.g.avante_login = vim.g.avante_login
 ---@field entries avante.ChatHistoryEntry[]
 ---@field memory avante.ChatMemory | nil
 ---@field filename string
+---@field system_prompt string | nil
 ---
 ---@class avante.ChatMemory
 ---@field content string
@@ -413,3 +414,11 @@ vim.g.avante_login = vim.g.avante_login
 ---@field content string
 ---@field uri string
 ---
+---@alias AvanteSlashCommandBuiltInName "clear" | "help" | "lines" | "reset" | "commit" | "new"
+---@alias AvanteSlashCommandCallback fun(self: avante.Sidebar, args: string, cb?: fun(args: string): nil): nil
+---@class AvanteSlashCommand
+---@field name AvanteSlashCommandBuiltInName | string
+---@field description string
+---@field details string
+---@field shorthelp? string
+---@field callback? AvanteSlashCommandCallback

lua/avante/ui/prompt_input.lua

@@ -95,7 +95,7 @@ function PromptInput:open()
 
   local bufnr = api.nvim_create_buf(false, true)
   self.bufnr = bufnr
-  vim.bo[bufnr].filetype = "AvanteInput"
+  vim.bo[bufnr].filetype = "AvantePromptInput"
   Utils.mark_as_sidebar_buffer(bufnr)
 
   local win_opts = vim.tbl_extend("force", {

lua/avante/utils/init.lua

@@ -1220,4 +1220,68 @@ function M.llm_tool_param_fields_to_json_schema(fields)
   return properties, required
 end
 
+---@return AvanteSlashCommand[]
+function M.get_commands()
+  local Config = require("avante.config")
+
+  ---@param items_ {name: string, description: string, shorthelp?: string}[]
+  ---@return string
+  local function get_help_text(items_)
+    local help_text = ""
+    for _, item in ipairs(items_) do
+      help_text = help_text .. "- " .. item.name .. ": " .. (item.shorthelp or item.description) .. "\n"
+    end
+    return help_text
+  end
+
+  local builtin_items = {
+    { description = "Show help message", name = "help" },
+    { description = "Clear chat history", name = "clear" },
+    { description = "Reset memory", name = "reset" },
+    { description = "New chat", name = "new" },
+    {
+      shorthelp = "Ask a question about specific lines",
+      description = "/lines <start>-<end> <question>",
+      name = "lines",
+    },
+    { description = "Commit the changes", name = "commit" },
+  }
+
+  ---@type {[AvanteSlashCommandBuiltInName]: AvanteSlashCommandCallback}
+  local builtin_cbs = {
+    help = function(sidebar, args, cb)
+      local help_text = get_help_text(builtin_items)
+      sidebar:update_content(help_text, { focus = false, scroll = false })
+      if cb then cb(args) end
+    end,
+    clear = function(sidebar, args, cb) sidebar:clear_history(args, cb) end,
+    reset = function(sidebar, args, cb) sidebar:reset_memory(args, cb) end,
+    new = function(sidebar, args, cb) sidebar:new_chat(args, cb) end,
+    lines = function(_, args, cb)
+      if cb then cb(args) end
+    end,
+    commit = function(_, _, cb)
+      local question = "Please commit the changes"
+      if cb then cb(question) end
+    end,
+  }
+
+  local builtin_commands = vim
+    .iter(builtin_items)
+    :map(
+      ---@param item AvanteSlashCommand
+      function(item)
+        return {
+          name = item.name,
+          description = item.description,
+          callback = builtin_cbs[item.name],
+          details = item.shorthelp and table.concat({ item.shorthelp, item.description }, "\n") or item.description,
+        }
+      end
+    )
+    :totable()
+
+  return vim.list_extend(builtin_commands, Config.slash_commands)
+end
+
 return M

lua/cmp_avante/commands.lua

@@ -1,40 +1,38 @@
 local api = vim.api
 
----@class commands_source : cmp.Source
+---@class CommandsSource : cmp.Source
----@field commands AvanteSlashCommand[]
+local CommandsSource = {}
----@field bufnr integer
+CommandsSource.__index = CommandsSource
-local commands_source = {}
-commands_source.__index = commands_source
 
----@param commands AvanteSlashCommand[]
+function CommandsSource:new()
----@param bufnr integer
+  local instance = setmetatable({}, CommandsSource)
-function commands_source:new(commands, bufnr)
-  local instance = setmetatable({}, commands_source)
-
-  instance.commands = commands
-  instance.bufnr = bufnr
 
   return instance
 end
 
-function commands_source:is_available() return api.nvim_get_current_buf() == self.bufnr end
+function CommandsSource:is_available() return vim.bo.filetype == "AvanteInput" end
 
-function commands_source.get_position_encoding_kind() return "utf-8" end
+function CommandsSource.get_position_encoding_kind() return "utf-8" end
 
-function commands_source:get_trigger_characters() return { "/" } end
+function CommandsSource:get_trigger_characters() return { "/" } end
 
-function commands_source:get_keyword_pattern() return [[\%(@\|#\|/\)\k*]] end
+function CommandsSource:get_keyword_pattern() return [[\%(@\|#\|/\)\k*]] end
 
-function commands_source:complete(_, callback)
+function CommandsSource:complete(_, callback)
+  local Utils = require("avante.utils")
   local kind = require("cmp").lsp.CompletionItemKind.Variable
+  local commands = Utils.get_commands()
 
   local items = {}
 
-  for _, command in ipairs(self.commands) do
+  for _, command in ipairs(commands) do
     table.insert(items, {
-      label = "/" .. command.command,
+      label = "/" .. command.name,
       kind = kind,
       detail = command.details,
+      data = {
+        name = command.name,
+      },
     })
   end
 
@@ -44,4 +42,20 @@ function commands_source:complete(_, callback)
   })
 end
 
-return commands_source
+function CommandsSource:execute(item, callback)
+  local Utils = require("avante.utils")
+  local commands = Utils.get_commands()
+  local command = vim.iter(commands):find(function(command) return command.name == item.data.name end)
+
+  if not command then return end
+
+  local sidebar = require("avante").get()
+  command.callback(sidebar, nil, function()
+    local content = table.concat(api.nvim_buf_get_lines(sidebar.input_container.bufnr, 0, -1, false), "\n")
+    content = content:gsub(item.label, "")
+    api.nvim_buf_set_lines(sidebar.input_container.bufnr, 0, -1, false, vim.split(content, "\n"))
+    callback()
+  end)
+end
+
+return CommandsSource

lua/cmp_avante/mentions.lua

@@ -1,36 +1,37 @@
 local api = vim.api
 
 ---@class mentions_source : cmp.Source
----@field mentions {description: string, command: AvanteMentions, details: string, shorthelp?: string, callback?: AvanteMentionCallback}[]
+---@field get_mentions fun(): {description: string, command: AvanteMentions, details: string, shorthelp?: string, callback?: AvanteMentionCallback}[]
----@field bufnr integer
+local MentionsSource = {}
-local mentions_source = {}
+MentionsSource.__index = MentionsSource
-mentions_source.__index = mentions_source
 
----@param mentions {description: string, command: AvanteMentions, details: string, shorthelp?: string, callback?: AvanteMentionCallback}[]
+---@param get_mentions fun(): {description: string, command: AvanteMentions, details: string, shorthelp?: string, callback?: AvanteMentionCallback}[]
----@param bufnr integer
+function MentionsSource:new(get_mentions)
-function mentions_source:new(mentions, bufnr)
+  local instance = setmetatable({}, MentionsSource)
-  local instance = setmetatable({}, mentions_source)
 
-  instance.mentions = mentions
+  instance.get_mentions = get_mentions
-  instance.bufnr = bufnr
 
   return instance
 end
 
-function mentions_source:is_available() return api.nvim_get_current_buf() == self.bufnr end
+function MentionsSource:is_available()
+  return vim.bo.filetype == "AvanteInput" or vim.bo.filetype == "AvantePromptInput"
+end
 
-function mentions_source.get_position_encoding_kind() return "utf-8" end
+function MentionsSource.get_position_encoding_kind() return "utf-8" end
 
-function mentions_source:get_trigger_characters() return { "@" } end
+function MentionsSource:get_trigger_characters() return { "@" } end
 
-function mentions_source:get_keyword_pattern() return [[\%(@\|#\|/\)\k*]] end
+function MentionsSource:get_keyword_pattern() return [[\%(@\|#\|/\)\k*]] end
 
-function mentions_source:complete(_, callback)
+function MentionsSource:complete(_, callback)
   local kind = require("cmp").lsp.CompletionItemKind.Variable
 
   local items = {}
 
-  for _, mention in ipairs(self.mentions) do
+  local mentions = self.get_mentions()
+
+  for _, mention in ipairs(mentions) do
     table.insert(items, {
       label = "@" .. mention.command .. " ",
       kind = kind,
@@ -46,22 +47,26 @@ end
 
 ---@param completion_item table
 ---@param callback fun(response: {behavior: number})
-function mentions_source:execute(completion_item, callback)
+function MentionsSource:execute(completion_item, callback)
   local current_line = api.nvim_get_current_line()
   local label = completion_item.label:match("^@(%S+)") -- Extract mention command without '@' and space
 
+  local mentions = self.get_mentions()
+
   -- Find the corresponding mention
   local selected_mention
-  for _, mention in ipairs(self.mentions) do
+  for _, mention in ipairs(mentions) do
     if mention.command == label then
       selected_mention = mention
       break
     end
   end
 
+  local sidebar = require("avante").get()
+
   -- Execute the mention's callback if it exists
   if selected_mention and type(selected_mention.callback) == "function" then
-    selected_mention.callback(selected_mention)
+    selected_mention.callback(sidebar)
     -- Get the current cursor position
     local row, col = unpack(api.nvim_win_get_cursor(0))
 
@@ -77,4 +82,4 @@ function mentions_source:execute(completion_item, callback)
   callback({ behavior = require("cmp").ConfirmBehavior.Insert })
 end
 
-return mentions_source
+return MentionsSource