Create context for LLM in neovim, save all your buffers as context.xml with F10

llm_context.lua

-----------------------------
-- Create context for the LLM
--
-- Save as: lua/llm_context.lua
--
-- How to use: require("llm_context").setup()
--
-- Reference:
-- Long context prompting tips
-- https://docs.anthropic.com/en/docs/build-with-claude/prompt-engineering/long-context-tips#example-quote-extraction
-----------------------------

local M = {}

-- Path for context file
local context_file = "context.xml"

-- Append content to context file
local function append_to_context(content)
  local file = io.open(context_file, "a")
  if file then
    file:write(content)
    file:close()
    -- vim.notify("Added to context file", vim.log.levels.INFO)
  else
    vim.notify("Failed to write context file", vim.log.levels.ERROR)
  end
end
-- Add buffer to context
local function add_buffer_to_context(bufnr)
  local bufname = vim.api.nvim_buf_get_name(bufnr)

  -- Skip if buffer has no name or is NvimTree
  if bufname == "" or bufname:match("NvimTree") then
    return false
  end

  -- Check if file exists on disk
  if vim.fn.filereadable(bufname) ~= 1 then
    return false
  end

  -- Check if file was already added by reading context.xml
  local current = io.open(context_file, "r")
  if current then
    local content = current:read("*all")
    current:close()
    if content:match(vim.fn.escape(bufname, "%-%.")) then
      return false
    end
  end

  -- current working directory
  local cwd = vim.fn.getcwd()

  -- remove cwd from bufname variable
  local relpath = vim.fn.substitute(bufname, cwd .. "/", "", "")

  local header = string.format("<document>\n<source>%s</source>\n<document_content>\n", relpath)
  local footer = "\n</document_content>\n</document>\n"

  local lines = vim.api.nvim_buf_get_lines(bufnr, 0, -1, false)
  local content = table.concat(lines, "\n")

  append_to_context(header .. content .. footer)
  return true
end

-------------------------
-- Add selection
-------------------------
vim.api.nvim_create_user_command("ContextAddSelection", function(opts)
  local start_line = vim.fn.line("'<") - 1
  local end_line = vim.fn.line("'>")

  local header = string.format("<document>\n<source>%s</source>\n<line_start>%d</line_start>\n<line_end>%d</line_end>\n<document_content>\n", vim.fn.expand("%:."), start_line + 1, end_line)
  local footer = "\n</document_content>\n</document>\n"

  local lines = vim.api.nvim_buf_get_lines(0, start_line, end_line, false)
  local content = table.concat(lines, "\n")

  append_to_context(header .. content .. footer)
  vim.notify("Added selection to context", vim.log.levels.INFO)
end, { range = true, bang = true })

-------------------------
-- Add current function
-------------------------
local function add_function()
  local current_pos = vim.fn.getpos(".")
  vim.cmd('normal "Zyas')

  local start_line = current_pos[2]
  local end_line = start_line + #vim.fn.split(vim.fn.getreg("Z"), "\n") - 1

  local header = string.format("<document>\n<source>%s</source>\n<line_start>%d</line_start>\n<line_end>%d</line_end>\n<document_content>\n", vim.fn.expand("%:."), start_line, end_line)
  local footer = "\n</document_content>\n</document>\n"

  local content = vim.fn.getreg("Z")
  append_to_context(header .. content .. footer)
  vim.notify("Added function to context", vim.log.levels.INFO)
end

------------------------
-- Add buffer to context
------------------------
local function add_buffer()
  if add_buffer_to_context(0) then
    vim.notify("Added file to context", vim.log.levels.INFO)
  end
end

------------------------
-- Add open buffers
------------------------
local function add_open_buffers()
  local buffers = vim.api.nvim_list_bufs()
  local added_count = 0

  for _, bufnr in ipairs(buffers) do
    if vim.api.nvim_buf_is_loaded(bufnr) then
      if add_buffer_to_context(bufnr) then
        added_count = added_count + 1
      end
    end
  end

  vim.notify(string.format("Added %d buffers to context", added_count), vim.log.levels.INFO)
end

---------------------------------------
-- Add directory structure to context
---------------------------------------
local function add_structure()
  if vim.fn.executable("tree") ~= 1 then
    vim.notify("'tree' command not found. Please install it first.", vim.log.levels.ERROR)
    return
  end

  local header = "<document>\n<source>directory_structure</source>\n<document_content>\n"
  local footer = "\n</document_content>\n</document>\n"

  local cmd = 'tree -L 3 -a --gitignore --noreport --dirsfirst -I ".git"'
  local lines = vim.fn.systemlist(cmd)
  local content = table.concat(lines, "\n")

  append_to_context(header .. content .. footer)
  vim.notify("Added directory structure to context", vim.log.levels.INFO)
end

------------------------------
-- Copy context to clipboard
------------------------------
local function dump_context()
  local file = io.open(context_file, "r")
  if file then
    local content = file:read("*all")
    file:close()
    vim.fn.setreg("+", "<documents>\n" .. content .. "</documents>")
    os.remove(context_file)
    vim.notify("Context copied to clipboard and file deleted", vim.log.levels.INFO)
  end
end
vim.api.nvim_create_user_command("ContextDump", dump_context, { bang = true })

--------------------------------
-- Add structure an open buffers
--------------------------------
local function add_editor()
  add_structure()
  add_open_buffers()
end

------------------------
-- Commands and mappings
------------------------
function M.setup()
  -- commands
  vim.api.nvim_create_user_command("ContextAddOpenBuffers", add_open_buffers, { bang = true })
  vim.api.nvim_create_user_command("ContextAddFunction", add_function, { bang = true })
  vim.api.nvim_create_user_command("ContextAddBuffer", add_buffer, { bang = true })
  vim.api.nvim_create_user_command("ContextAddStructure", add_structure, { bang = true })
  vim.api.nvim_create_user_command("ContextAddEditor", add_editor, { bang = true })
  -- mappings
  vim.keymap.set("v", "<F9>", ":<C-u>ContextAddSelection<CR>", { noremap = true, silent = true })
  vim.keymap.set("n", "<F9>", ":<C-u>ContextAddBuffer<CR>", { noremap = true, silent = true })
  vim.keymap.set("n", "<F10>", ":ContextAddEditor<CR>", { noremap = true, silent = true })
  vim.keymap.set("n", "<F12>", ":ContextDump<CR>", { noremap = true, silent = true })
end

return M