App Tools

How to add an application tool to llm.nvim?

Important

Currently, llm.nvim defines some tool templates by default, which you can easily call and modify prompts and some of its options.

{
  "Kurama622/llm.nvim",
  dependencies = { "nvim-lua/plenary.nvim", "MunifTanjim/nui.nvim" },
  cmd = { "LLMSesionToggle", "LLMSelectedTextHandler", "LLMAppHandler" },
  config = function()
    local tools = require("llm.common.tools")
    require("llm").setup({
      app_handler = {
        OptimizeCode = {
          handler = tools.side_by_side_handler,
        },

        -- modify the prompt and options
        TestCode = {
          handler = tools.side_by_side_handler,
          prompt = "Write a test for the following code, only return the test code:",
          opts = {
            right = {
              title = " Result ",
            },
          },
        },
        OptimCompare = {
          handler = tools.action_handler,
        },
        Translate = {
          handler = tools.qa_handler,
        },
      },
    })
  end,
  keys = {
    { "<leader>tc", mode = "n", "<cmd>LLMAppHandler TestCode<cr>" },
    { "<leader>t", mode = "x", "<cmd>LLMSelectedTextHandler 英译汉<cr>" },
    { "<leader>at", mode = "n", "<cmd>LLMAppHandler Translate<cr>" },
    { "<leader>ao", mode = "x", "<cmd>LLMAppHandler OptimCompare<cr>" },
  },
}

template	opts	example
action_handler	func code_hl separator_hl border win_options buftype spell number wrap linebreak	options = { func = CompareAction, code_hl = { fg = "#6aa84f", bg = "NONE" }, separator_hl = { fg = "#6aa84f", bg = "#333333" }, border = "solid", win_options = { winblend = 0, winhighlight = "Normal:Normal" }, buftype = "nofile", spell = false, number = true, wrap = true, linebreak = false, }
side_by_side_handler	left right buftype spell number wrap linebreak	options = { left = { title = " Source ", }, right = { title = " Preview ", }, buftype = "nofile", spell = false, number = true, wrap = true, linebreak = false, }
qa_handler	query buftype spell number wrap linebreak	options = { query = { title = " 󰊿 Trans ", hl = { link = "CurSearch" }, }, buftype = "nofile", spell = false, number = false, wrap = true, linebreak = false, }

You can also customize your tools from scratch:

1. Define your handler function

2. Insert the handler function into app_handler

{
  "Kurama622/llm.nvim",
  dependencies = { "nvim-lua/plenary.nvim", "MunifTanjim/nui.nvim" },
  cmd = { "LLMSesionToggle", "LLMSelectedTextHandler", "LLMAppHandler" },
  config = function()
    require("llm").setup({
      url = "https://open.bigmodel.cn/api/paas/v4/chat/completions",
      model = "glm-4-flash",
      max_tokens = 4095,

      prompt = "",

      prefix = {
        user = { text = "😃 ", hl = "Title" },        -- 
        assistant = { text = "  ", hl = "Added" },
      },

      save_session = true,
      max_history = 15,

      app_handler = {
        ToolName1 = handler1,
        ToolName2 = handler2,
      },
    })
  end
}

Note

All functions in https://github.com/Kurama622/llm.nvim/blob/main/lua/llm/common/func.lua can be called in the form of F.xxx

Some code snippets placed at the beginning.

Tip

Currently these functions have been integrated into https://github.com/Kurama622/llm.nvim/blob/main/lua/llm/common/func.lua and no longer need to be defined.

local Popup = require("nui.popup")
local Layout = require("nui.layout")
local NuiText = require("nui.text")

local function InsertTextLine(bufnr, linenr, text)
  vim.api.nvim_buf_set_lines(bufnr, linenr, linenr, false, { text })
end

local function ReplaceTextLine(bufnr, linenr, text)
  vim.api.nvim_buf_set_lines(bufnr, linenr, linenr + 1, false, { text })
end

local function RemoveTextLines(bufnr, start_linenr, end_linenr)
  vim.api.nvim_buf_set_lines(bufnr, start_linenr, end_linenr, false, {})
end

-- create a popup
local function CreatePopup(text, focusable, opts)
  local options = {
    focusable = focusable,
    border = { style = "rounded", text = { top = text, top_align = "center" } },
  }
  options = vim.tbl_deep_extend("force", options, opts or {})

  return Popup(options)
end

-- create your layout
local function CreateLayout(_width, _height, boxes, opts)
  local options = {
    relative = "editor",
    position = "50%",
    size = {
      width = _width,
      height = _height,
    },
  }
  options = vim.tbl_deep_extend("force", options, opts or {})
  return Layout(options, boxes)
end

-- Set the properties of the pop-up window.
local function SetBoxOpts(box_list, opts)
  for i, v in ipairs(box_list) do
    vim.api.nvim_set_option_value("filetype", opts.filetype[i], { buf = v.bufnr })
    vim.api.nvim_set_option_value("buftype", opts.buftype, { buf = v.bufnr })
    vim.api.nvim_set_option_value("spell", opts.spell, { win = v.winid })
    vim.api.nvim_set_option_value("wrap", opts.wrap, { win = v.winid })
    vim.api.nvim_set_option_value("linebreak", opts.linebreak, { win = v.winid })
    vim.api.nvim_set_option_value("number", opts.number, { win = v.winid })
  end
end

Create a tool to help optimize your code.

Tip

Currently these functions have been integrated into https://github.com/Kurama622/llm.nvim/blob/main/lua/llm/common/tools.lua and no longer need to be defined.

local optimize_code_handler = function(name, F, state, streaming)
  local ft = vim.bo.filetype
  local prompt = [[优化代码, 修改语法错误, 让代码更简洁, 增强可复用性，
            你要像copliot那样，直接给出代码内容, 不要使用代码块或其他标签包裹!

            下面是一个例子，假设我们需要优化下面这段代码:
            void test() {
             return 0
            }

            输出格式应该为：
            int test() {
              return 0;
            }

            请按照格式，帮我优化这段代码：
            ]]

  local source_content = F.GetVisualSelection()

  local source_box = CreatePopup(" Source ", false)
  local preview_box = CreatePopup(" Preview ", true, { enter = true })

  local layout = CreateLayout(
    "80%",
    "55%",
    Layout.Box({
      Layout.Box(source_box, { size = "50%" }),
      Layout.Box(preview_box, { size = "50%" }),
    }, { dir = "row" })
  )

  layout:mount()

  SetBoxOpts({ source_box, preview_box }, {
    filetype = { ft, ft },
    buftype = "nofile",
    spell = false,
    number = true,
    wrap = true,
    linebreak = false,
  })

  state.popwin = source_box
  F.WriteContent(source_box.bufnr, source_box.winid, source_content)

  state.app["session"][name] = {}
  table.insert(state.app.session[name], { role = "user", content = prompt .. "\n" .. source_content })

  state.popwin = preview_box
  local worker = streaming(preview_box.bufnr, preview_box.winid, state.app.session[name])

  preview_box:map("n", "<C-c>", function()
    if worker.job then
      worker.job:shutdown()
      worker.job = nil
    end
  end)

  preview_box:map("n", { "<esc>", "N", "n" }, function()
    if worker.job then
      worker.job:shutdown()
      print("Suspend output...")
      vim.wait(200, function() end)
      worker.job = nil
    end
    layout:unmount()
  end)

  preview_box:map("n", { "Y", "y" }, function()
    vim.api.nvim_command("normal! ggVGky")
    layout:unmount()
  end)
end

Create a tool to help optimize your code and show the result in source file.

Tip

Currently these functions have been integrated into https://github.com/Kurama622/llm.nvim/blob/main/lua/llm/common/tools.lua and no longer need to be defined.

local CompareAction = function(bufnr, start_str, end_str, mark_id, extmark,
  extmark_opts, space_text, start_line, end_line, codeln, offset, ostr)
  local pattern = string.format("%s(.-)%s", start_str, end_str)
  local res = ostr:match(pattern)
  if res == nil then
    print("The code block format is incorrect, please manually copy the generated code.")
    return codeln
  end

  vim.api.nvim_set_hl(0, "LLMSuggestCode", { fg = "#6aa84f", bg = "NONE" })
  vim.api.nvim_set_hl(0, "LLMSeparator", { fg = "#6aa84f", bg = "#333333" })

  for _, v in ipairs({ "raw", "separator", "llm" }) do
    extmark[v] = vim.api.nvim_create_namespace(v)
    local text = v == "raw" and "<<<<<<< " .. v .. space_text
      or v == "separator" and "======= " .. space_text
      or ">>>>>>> " .. v .. space_text
    extmark_opts[v] = {
      virt_text = { { text, "LLMSeparator" } },
      virt_text_pos = "overlay",
    }
  end

  extmark["code"] = vim.api.nvim_create_namespace("code")
  extmark_opts["code"] = {}
  mark_id["code"] = {}

  if offset ~= 0 then
    -- create line to display raw separator virtual text
    InsertTextLine(bufnr, 0, "")
  end

  mark_id["raw"] = vim.api.nvim_buf_set_extmark(bufnr, extmark.raw, start_line - 2 + offset, 0, extmark_opts.raw)

  -- create line to display the separator virtual text
  InsertTextLine(bufnr, end_line + offset, "")
  mark_id["separator"] =
    vim.api.nvim_buf_set_extmark(bufnr, extmark.separator, end_line + offset, 0, extmark_opts.separator)

  for l in res:gmatch("[^\r\n]+") do
    -- create line to display the code suggested by the LLM
    InsertTextLine(bufnr, end_line + codeln + 1 + offset, "")
    extmark_opts.code[codeln] = { virt_text = { { l, "LLMSuggestCode" } }, virt_text_pos = "overlay" }
    mark_id.code[codeln] =
      vim.api.nvim_buf_set_extmark(bufnr, extmark.code, end_line + codeln + 1 + offset, 0, extmark_opts.code[codeln])
    codeln = codeln + 1
  end

  -- create line to display LLM separator virtual text
  InsertTextLine(bufnr, end_line + codeln + 1 + offset, "")
  mark_id["llm"] = vim.api.nvim_buf_set_extmark(bufnr, extmark.llm, end_line + codeln + 1 + offset, 0, extmark_opts.llm)
  return codeln
end

local optim_compare_action_handler = function(name, F, state, streaming)
  local prompt = [[Optimize the code, correct syntax errors, make the code more concise, and enhance reusability.

Provide optimization ideas and the complete code after optimization. Mark the output code block with # BEGINCODE and # ENDCODE.

The indentation of the optimized code should remain consistent with the original code. Here is an example:

The original code is:
<space><space><space><space>def func(a, b)
<space><space><space><space><space><space><space><space>return a + b

Optimization ideas:
1. The function name `func` is not clear. Based on the context, it is determined that this function is meant to implement the functionality of adding two numbers, so the function name is changed to `add`.
2. There is a syntax issue in the function definition; it should end with a colon. It should be `def add(a, b):`.

Since the original code is indented by N spaces, the optimized code is also indented by N spaces.

The optimized code is:

```<language>
# BEGINCODE
<space><space><space><space>def add(a, b):
<space><space><space><space><space><space><space><space>return a + b
# ENDCODE
```

Please optimize this code according to the format, and respond in Chinese.]]
  local start_line, end_line = F.GetVisualSelectionRange()
  local bufnr = vim.api.nvim_get_current_buf()
  local source_content = F.GetVisualSelection()

  local preview_box = Popup({ enter = true, border = "solid",
    win_options = { winblend = 0, winhighlight = "Normal:Normal" } })

  local layout = CreateLayout( "30%", "98%",
    Layout.Box({ Layout.Box(preview_box, { size = "100%" }) }, { dir = "row" }),
    { position = { row = "50%", col = "100%" } })

  layout:mount()

  local mark_id = {}
  local extmark = {}
  local extmark_opts = {}
  local space_text = string.rep(" ", vim.o.columns - 7)
  local start_str = "# BEGINCODE"
  local end_str = "# ENDCODE"
  local codeln = 0
  local offset = start_line == 1 and 1 or 0

  SetBoxOpts({ preview_box }, {
    filetype = { "markdown" },
    buftype = "nofile",
    spell = false,
    number = true,
    wrap = true,
    linebreak = false,
  })

  state.app["session"][name] = {}
  table.insert(state.app.session[name], { role = "user", content = prompt .. "\n" .. source_content })

  state.popwin = preview_box
  local worker = streaming(
    preview_box.bufnr,
    preview_box.winid,
    state.app.session[name],
    nil, -- curl args
    nil, -- streaming handler
    nil, -- stdout handler
    nil, -- stderr handler
    function(ostr) -- exit handler
      codeln = CompareAction(bufnr, start_str, end_str, mark_id, extmark,
        extmark_opts, space_text, start_line, end_line, codeln, offset, ostr)
    end
  )

  preview_box:map("n", "<C-c>", function()
    if worker.job then
      worker.job:shutdown()
      worker.job = nil
    end
  end)

  preview_box:map("n", { "<esc>", "N", "n" }, function()
    if worker.job then
      worker.job:shutdown()
      print("Suspend output...")
      vim.wait(200, function() end)
      worker.job = nil
    end
    if codeln ~= 0 then
      vim.api.nvim_buf_del_extmark(bufnr, extmark.raw, mark_id.raw)
      vim.api.nvim_buf_del_extmark(bufnr, extmark.separator, mark_id.separator)
      vim.api.nvim_buf_del_extmark(bufnr, extmark.llm, mark_id.llm)
      for i = 0, codeln - 1 do
        vim.api.nvim_buf_del_extmark(bufnr, extmark.code, mark_id.code[i])
      end

      -- remove the line created to display the code suggested by LLM.
      RemoveTextLines(bufnr, end_line + offset, end_line + codeln + 2 + offset)
      if offset ~= 0 then
        -- remove the line created to display the raw separator.
        RemoveTextLines(bufnr, 0, 1)
      end
    end
    layout:unmount()
  end)

  preview_box:map("n", { "Y", "y" }, function()
    if codeln ~= 0 then
      vim.api.nvim_buf_del_extmark(bufnr, extmark.raw, mark_id.raw)
      vim.api.nvim_buf_del_extmark(bufnr, extmark.separator, mark_id.separator)
      vim.api.nvim_buf_del_extmark(bufnr, extmark.llm, mark_id.llm)

      -- remove the line created to display the LLM separator.
      RemoveTextLines(bufnr, end_line + codeln + 1 + offset, end_line + codeln + 2 + offset)
      -- remove raw code
      RemoveTextLines(bufnr, start_line - 1, end_line + 1 + offset)

      for i = 0, codeln - 1 do
        vim.api.nvim_buf_del_extmark(bufnr, extmark.code, mark_id.code[i])
      end

      for i = 0, codeln - 1 do
        -- Write the code suggested by the LLM.
        ReplaceTextLine(bufnr, start_line - 1 + i, extmark_opts.code[i].virt_text[1][1])
      end
    end
    layout:unmount()
  end)
end

Create a translator tool.

Tip

Currently these functions have been integrated into https://github.com/Kurama622/llm.nvim/blob/main/lua/llm/common/tools.lua and no longer need to be defined.

local translate_handler = function(name, _, state, streaming)
  local prompt = [[请帮我把这段话翻译成英语, 直接给出翻译结果: ]]

  local input_box = Popup({
    enter = true,
    border = {
      style = "solid",
      text = {
        top = NuiText(" 󰊿 Trans ", "CurSearch"),
        top_align = "center",
      },
    },
  })

  local separator = Popup({
    border = { style = "none" },
    enter = false,
    focusable = false,
    win_options = { winblend = 0, winhighlight = "Normal:Normal" },
  })

  local preview_box = Popup({
    focusable = true,
    border = { style = "solid", text = { top = "", top_align = "center" } },
  })

  local layout = CreateLayout(
    "60%",
    "55%",
    Layout.Box({
      Layout.Box(input_box, { size = "15%" }),
      Layout.Box(separator, { size = "5%" }),
      Layout.Box(preview_box, { size = "80%" }),
    }, { dir = "col" })
  )

  layout:mount()
  vim.api.nvim_command("startinsert")

  SetBoxOpts({ preview_box }, {
    filetype = { "markdown", "markdown" },
    buftype = "nofile",
    spell = false,
    number = false,
    wrap = true,
    linebreak = false,
  })

  local worker = { job = nil }

  state.app["session"][name] = {}
  input_box:map("n", "<enter>", function()
    -- clear preview_box content [optional]
    vim.api.nvim_buf_set_lines(preview_box.bufnr, 0, -1, false, {})

    local input_table = vim.api.nvim_buf_get_lines(input_box.bufnr, 0, -1, true)
    local input = table.concat(input_table, "\n")

    -- clear input_box content
    vim.api.nvim_buf_set_lines(input_box.bufnr, 0, -1, false, {})
    if input ~= "" then
      table.insert(state.app.session[name], { role = "user", content = prompt .. "\n" .. input })
      state.popwin = preview_box
      worker = streaming(preview_box.bufnr, preview_box.winid, state.app.session[name])
    end
  end)

  input_box:map("n", { "J", "K" }, function()
    vim.api.nvim_set_current_win(preview_box.winid)
  end)
  preview_box:map("n", { "J", "K" }, function()
    vim.api.nvim_set_current_win(input_box.winid)
  end)

  for _, v in ipairs({ input_box, preview_box }) do
    v:map("n", { "<esc>", "N", "n" }, function()
      if worker.job then
        worker.job:shutdown()
        print("Suspend output...")
        vim.wait(200, function() end)
        worker.job = nil
      end
      layout:unmount()
    end)

    v:map("n", { "Y", "y" }, function()
      vim.api.nvim_set_current_win(preview_box.winid)
      vim.api.nvim_command("normal! ggVGky")
      layout:unmount()
    end)
  end
end