Skip to content

Modify LLM

Kurama edited this page Sep 24, 2024 · 6 revisions

How to use custom LLM models?

  1. Add the requested URL.
  2. Specify the model you will be using.
  3. Customize the streaming processing function (used for parsing the model output).

Example

kimi

local kimi_handler = function(chunk, line, output, bufnr, winid, F)
  if not chunk then
    return output
  end
  local tail = chunk:sub(-1, -1)
  if tail:sub(1, 1) ~= "}" then
    line = line .. chunk
  else
    line = line .. chunk

    local start_idx = line:find("data: ", 1, true)
    local end_idx = line:find("}]", 1, true)
    local json_str = nil

    while start_idx ~= nil and end_idx ~= nil do
      if start_idx < end_idx then
        json_str = line:sub(7, end_idx + 1) .. "}"
      end
      local data = vim.fn.json_decode(json_str)
      if not data.choices[1].delta.content then
        break
      end

      output = output .. data.choices[1].delta.content
      F.WriteContent(bufnr, winid, data.choices[1].delta.content)

      if end_idx + 2 > #line then
        line = ""
        break
      else
        line = line:sub(end_idx + 2)
      end
      start_idx = line:find("data: ", 1, true)
      end_idx = line:find("}]", 1, true)
    end
  end
  return output
end

return {
  {
    "Kurama622/llm.nvim",
    dependencies = { "nvim-lua/plenary.nvim", "MunifTanjim/nui.nvim" },
    cmd = { "LLMSesionToggle", "LLMSelectedTextHandler", "LLMAppHandler" },
    config = function()
      require("llm").setup({
        -- kimi
        url = "https://api.moonshot.cn/v1/chat/completions",
        model = "moonshot-v1-8k", -- "moonshot-v1-8k", "moonshot-v1-32k", "moonshot-v1-128k"
        streaming_handler = kimi_handler,

        max_tokens = 4096,
        -- temperature = 0.3,
        -- top_p = 0.7,

        prompt = [[]],

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

        save_session = true,
        max_history = 15,
        max_history_name_length = 12,
      })
    end,
    keys = {
      { "<leader>ac", mode = "n", "<cmd>LLMSessionToggle<cr>" },
      { "<leader>ae", mode = "v", "<cmd>LLMSelectedTextHandler 请解释下面这段代码<cr>" },
      { "<leader>t", mode = "x", "<cmd>LLMSelectedTextHandler 英译汉<cr>" },
    },
  },
}

Github Models

local github_model_handler = function(chunk, line, output, bufnr, winid, F)
  if not chunk then
    return output
  end
  local tail = chunk:sub(-1, -1)
  if tail:sub(1, 1) ~= "}" then
    line = line .. chunk
  else
    line = line .. chunk

    local start_idx = line:find("data: ", 1, true)
    local end_idx = line:find("}]", 1, true)
    local json_str = nil

    while start_idx ~= nil and end_idx ~= nil do
      if start_idx < end_idx then
        json_str = line:sub(7, end_idx + 1) .. "}"
      end
      local data = vim.fn.json_decode(json_str)
      if not data.choices[1].delta.content then
        break
      end

      output = output .. data.choices[1].delta.content
      F.WriteContent(bufnr, winid, data.choices[1].delta.content)

      if end_idx + 2 > #line then
        line = ""
        break
      else
        line = line:sub(end_idx + 2)
      end
      start_idx = line:find("data: ", 1, true)
      end_idx = line:find("}]", 1, true)
    end
  end
  return output
end

return {
  {
    "Kurama622/llm.nvim",
    dependencies = { "nvim-lua/plenary.nvim", "MunifTanjim/nui.nvim" },
    cmd = { "LLMSesionToggle", "LLMSelectedTextHandler", "LLMAppHandler" },
    config = function()
      require("llm").setup({
        -- Github models
        url = "https://models.inference.ai.azure.com/chat/completions",
        model = "gpt-4o",
        streaming_handler = github_model_handler,

        max_tokens = 4096,
        -- temperature = 0.3,
        -- top_p = 0.7,

        prompt = [[]],

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

        save_session = true,
        max_history = 15,
        max_history_name_length = 12,
      })
    end,
    keys = {
      { "<leader>ac", mode = "n", "<cmd>LLMSessionToggle<cr>" },
      { "<leader>ae", mode = "v", "<cmd>LLMSelectedTextHandler 请解释下面这段代码<cr>" },
      { "<leader>t", mode = "x", "<cmd>LLMSelectedTextHandler 英译汉<cr>" },
    },
  },
}

If you are a paid user of azure, you only need to modify the `url`.
Clone this wiki locally