🤔 Vineeth's Thoughts

Search

SearchSearch

Recent Posts

See 12 more →

Latios Devlog #1 A Naive Copilot

Sep 08, 202414 min read

This post is in relation to my [[latios|latios.nvim] project, which can be found on GitHub.

At the time of writing this post I’ve built a naive AI copilot as a plugin for neovim using lua. I say naive because it…

  • Is slow
  • Doesn’t really have the best completions
  • Isn’t super easily hackable
  • Not rigorously tested for backwards compatibility

But what we do have is a copilot that

  • shows inline completions that can be Tab completed
  • builds a context using treesitter, lsp, and several other sources
  • requests completions from claude sonnet 3.5

A decent starting point. My process for building this was to analyze a few other existing copilot solutions in neovim using repo-to-prompt and claude. The artifacts feature had just released when I started development and a lot of those were incorporated into the initial designs.

Note

If you want to follow along in my repo I’m writing this at the point of commit 199d4ac

A lot of the code is redundant and definitely needs some cleaning up before it’s presentable, but I’ll do my best to only cover the parts that are actually used. Some aspects of the code are related to neovim plugins in general while others are focused on the mechanics of a copilot. I’ll make an effort to delineate between the two in my description

Structure

Neovim Plugin Structure

The structure of the repo itself is very similar to how today.nvim is structured.

latios/
├─ doc/
│  ├─ latios.txt
├─ lua/latios/
│  ├─ init.lua
│  ├─ ...
│  ├─ ...
├─ plugin/
│  ├─ latios.lua
├─ README.md

The doc folder contains the documentation that will be displayed when running help on a plugin in neovim. The latios.txt file has some really bare bones content right now.

The plugin folder itself just has a single file that is used to load the lua module and run some simple setup code.

if vim.g.loaded_latios then
  return
end
 
vim.g.loaded_latios = true
 
-- Set up any global variables or initial configurations here
vim.g.latios_enabled = true
 
vim.api.nvim_create_user_command('Latios', function(opts)
  require('latios.commands').handle_command(opts.args)
end, {
  nargs = '*',
  complete = function(_, line)
    local commands = { 'enable', 'disable', 'toggle' }
    return vim.tbl_filter(function(cmd)
      return cmd:match('^' .. line)
    end, commands)
  end
})
.
.
.
-- Load the main module
require('latios').setup()

We use a global variable to determine if the plugin is enabled at all. If it is, then we register a user command to take in other inputs. This makes it so users can manually toggle things while in a neovim session using syntax like Latios <cmd>. There are a few commands there right now just for toggling the functionality.

Finally, we load the lua/ modules and run the entrypoint setup function.

This structure is so we can take advantage of lazy-loading. Only the setup code is run when a neovim buffer is opened. The rest of the functionality of the plugin doesn’t load until you run something that will cause a require() call.

Everything in the lua/ folder is requirable by the latio.lua file so by simply requiring the name of the plugin it will load the entrypoint init.lua file and then run the setup() function within.

Co-pilot Structure

The copilot’s logic is all contained within the lua/latios folder. Zooming in on this folder we have

latios/
├─ plugin/latios/
│  ├─ context.lua
│  ├─ init.lua
│  ├─ server.lua
│  ├─ display.lua

There are some other files in the codebase, but as mentioned earlier they are not used and contain some redundant code that was generated by claude.

To highlight how the functionality is split between these modules

  • init.lua sets up the autocommands that will trigger LLM completions
  • context.lua manages deriving the context to send to an LLM so it knows what to generate
  • server.lua manages sending the context to the LLM to get the completion
  • display.lua manages displaying the potential completion to the screen and writing it to the buffer if accepted

Zooming in even more on each of these modules…

init.lua

function M.setup(opts)
  require('latios.config').setup(opts)
 
  vim.api.nvim_create_autocmd({ "InsertEnter", "CursorMovedI", "CompleteChanged" }, {
    callback = function()
      if vim.g.latios_enabled and not utils.is_telescope_buffer() then
        server.debounced_request_completion(function(completion)
          if completion then
            display.show_completion(completion)
          end
        end)
      end
    end,
  })
 
  vim.api.nvim_create_autocmd({ 'InsertLeave', 'BufUnload' }, {
    callback = function()
      is_insert_mode = false
      if vim.g.latios_enabled and not utils.is_telescope_buffer() then
        server.cancel_ongoing_requests()
        display.clear_completion()
      end
    end,
  })
end

The main content of this file is the setup function that was called from plugin/latios.lua. It first loads configuration variables that would be setup via a package manager such as Lazy.nvim. Then it creates 2 autocommands.

The first one will trigger a new completion when the user enters insert mode, the user moves the cursor in insert mode, and when the completion options in a completion plugin like nvim-cmp. If a completion request yields a result, then the plugin will display it inline in the buffer.

The second one will clear the display and cancel any in-flight requests when the user exits insert mode. This is to prevent any excess requests and to remove stale completion results from the screen.

context.lua

There are a few different functions in this module that work to build context from multiple different sources within the editor. I’ll cover the main two.

get_treesitter_context gets the node the cursor is at and returns the tree up to that point.

function M.get_treesitter_context()
  -- Check if there's a parser available for the current buffer
  if not vim.treesitter.language.get_lang(vim.bo.filetype) then
    return nil
  end
 
  local context = {}
 
  -- Ensure we have a parser for the current buffer
  local parser = vim.treesitter.get_parser(0)
  if not parser then
    return nil
  end
 
  local tree = parser:parse()[1]
  local root = tree:root()
 
  -- Get the node at the cursor
  local cursor = vim.api.nvim_win_get_cursor(0)
  local node = root:named_descendant_for_range(cursor[1] - 1, cursor[2], cursor[1] - 1, cursor[2])
 
  -- Traverse up the tree to get context
  while node do
    table.insert(context, { type = node:type(), text = vim.treesitter.get_node_text(node, 0) })
    node = node:parent()
  end
 
  return context
end

get_lsp_context which uses the built in lsp to get any diagnostics, the symbol under the cursor, and the function signature if the cursor is within a function.

function M.get_lsp_context()
  local context = {}
 
  -- Get current buffer diagnostics
  context.diagnostics = vim.lsp.diagnostic.get_line_diagnostics()
 
  -- Get symbol under cursor
  local params = vim.lsp.util.make_position_params()
  local result = vim.lsp.buf_request_sync(0, 'textDocument/hover', params, 1000)
  if result and result[1] then
    context.hover_info = result[1].result
  end
 
  -- Get function signature if inside a function call
  local signature = vim.lsp.buf_request_sync(0, 'textDocument/signatureHelp', params, 1000)
  if signature and signature[1] then
    context.signature_help = signature[1].result
  end
 
  return context
end

The remaining functions will use those two sources along with the contents of the buffer and the filetype to build a table that is used as the context. There is also a reduced context function that attempts to truncate the lines of the buffer that are returned and another function to cache the results. Both of these latter functions were some early attempts for optimizing performance, but are largely un-tested.

server.lua

The two main functions in this module are construct_prompt which takes the context and makes a proper prompt for the language model and request_anthropic_completion which uses that prompt and sends a request to the Anthropic API for a completion

First construct_prompt

local function construct_prompt(full_context)
  -- return profiler.profile("construct_prompt", function()
  -- System prompt
  local system_prompt = [[
  You are an AI programming assistant specialized in providing code completions
  and explanations. Your responses should be concise, relevant, and tailored to
  the specific programming context provided. Do not provide any explanation or
  styling only the code itself. Do not talk at all. Only output valid code. Do
  not provide any backticks that surround the code. Never ever output backticks
  like this ```.]]
 
  -- local prompt = "You are an AI programming assistant. "
  local prompt = string.format("I'm working on a %s file: %s. ",
    full_context.filetype,
    full_context.file_path or "unknown path"
  )
 
  prompt = prompt .. string.format("My cursor is at position %d:%d. ",
    full_context.cursor_position[1],
    full_context.cursor_position[2]
  )
  -- "My cursor is at position " .. full_context.cursor_position[1] .. ":" .. full_context.cursor_position[2] .. ". "
 
  -- Add LSP context
  if full_context.lsp.hover_info then
    prompt = prompt .. "The symbol under cursor is: " .. tostring(full_context.lsp.hover_info) .. ". "
  end
  if full_context.lsp.signature_help then
    prompt = prompt .. "The current function signature is: " .. tostring(full_context.lsp.signature_help) .. ". "
  end
 
  -- Add Tree-sitter context
  if full_context.treesitter then
    prompt = prompt .. "Syntactic context: "
    for _, node in ipairs(full_context.treesitter) do
      prompt = prompt .. string.format("%s (%s), ", node.type, node.text)
      -- prompt = prompt .. node.type .. " (" .. node.text .. "), "
    end
  end
 
  -- Add surrounding code context
  local lines = full_context.buffer_content
  local cursor_line = full_context.cursor_position[1]
  local context_size = math.min(10, math.floor(#lines / 2)) -- Adjust based on file size
  local start_line = math.max(1, cursor_line - context_size)
  local end_line = math.min(#lines, cursor_line + context_size)
 
  prompt = prompt .. "Surrounding code:\n"
 
  for i = start_line, end_line do
    if i == cursor_line then
      prompt = prompt .. "> " .. lines[i] .. "\n" -- Highlight current line
    else
      prompt = prompt .. lines[i] .. "\n"
    end
  end
 
  -- Include function/class context if available
  if full_context.current_function then
    prompt = prompt .. "Current function:\n" .. full_context.current_function .. "\n"
  end
 
  -- Add language-specific context
  if full_context.language_specific then
    prompt = prompt .. "Language-specific context: " .. full_context.language_specific .. "\n"
  end
 
  prompt = prompt ..
      "Please provide a completion for the current position. If there is not a relevant completion output nil"
 
  return {
    system_prompt = system_prompt,
    messages = {
      { role = "user", content = prompt }
    }
  }
  -- end)(full_context)
end

The function essentially takes the table returned from context.lua and tries to put it more into prose for the language model. It also structures it as a set of messages so it matches the chat model API specs.

Next in request_anthropic_completion

local function request_anthropic_completion(prompt_data, callback)
  -- return profiler.profile("request_anthropic_completion", function()
  local request_body = vim.fn.json_encode({
    system = prompt_data.system_prompt,
    messages = prompt_data.messages,
    -- model = "claude-3-haiku-20240307",
    model = "claude-3-5-sonnet-20240620",
    max_tokens = 100,
    stop_sequences = { "\n\nHuman:" },
    temperature = 0.8,
  })
 
  local stdout = {}
  local stderr = {}
 
  local job_id = vim.fn.jobstart({
    'curl',
    '-s',
    '-H', 'Content-Type: application/json',
    '-H', 'X-API-Key: ' .. config.options.api_key,
    '-H', 'Anthropic-Version: 2023-06-01',
    '-d', request_body,
    'https://api.anthropic.com/v1/messages'
  }, {
    stdout_buffered = true,
    stderr_buffered = true,
    on_stdout = function(_, data)
      if data then
        vim.list_extend(stdout, data)
      end
    end,
    on_stderr = function(_, data)
      if data then
        vim.list_extend(stderr, data)
      end
    end,
    on_exit = function(_, exit_code)
      if exit_code == 0 and #stdout > 0 then
        local response = vim.fn.json_decode(table.concat(stdout))
        local completion = response.content[1].text
        bug.debug_info(completion)
        callback(completion)
      else
        print("Error from Anthropic API:", table.concat(stderr, '\n'))
        bug.debug_error(stderr)
        callback(nil)
      end
    end,
  })
 
  return job_id
end

This function launches an http request using curl and then has some callback functions depending on the output. The request is initiated via a jobstart command. This is to manage some asynchronous behavior and make it easier to cancel requests. The rest of the file is related to functions helping to manage the asynchronous nature of the request.

display.lua

This is the last big module in the codebase. The main purpose is to be able to display inline completions the same way apps like GitHub Copilot or Cursor, rather than just using the completions popup.

This is achieved by using neovim’s Extended Marks feature. The 3 main functions are used to show, clear, and accept completions returned by server.lua.

function M.show_completion(completion)
  if completion == 'nil' then return end
  local bufnr = vim.api.nvim_get_current_buf()
  local ns_id = vim.api.nvim_create_namespace('latios')
  local line, col = unpack(vim.api.nvim_win_get_cursor(0))
  line = line - 1 -- API uses 0-based line numbers
 
  -- Clear the previous extmark if it exists
  if current_extmark_id then
    vim.api.nvim_buf_del_extmark(bufnr, ns_id, current_extmark_id)
  end
 
  -- Split the completion into lines
  local lines = vim.split(completion, '\n', true)
 
  -- Remove the last line if it's empty
  if lines[#lines] == '' then
    table.remove(lines, #lines)
  end
  local opts = {
    virt_text = { { lines[1], 'Comment' } },
    virt_text_pos = 'overlay',
    hl_mode = 'combine',
 
  }
 
  -- If there are additional lines, add them as virt_lines
  if #lines > 1 then
    opts.virt_lines = {}
    for i = 2, #lines do
      table.insert(opts.virt_lines, { { lines[i], 'Comment' } })
    end
  end
 
  -- Set the new extmark and store its ID
  current_extmark_id = vim.api.nvim_buf_set_extmark(bufnr, ns_id, line, col, opts)
end

The function splits the completion into lines and displays them using the same Highlight as comments. In the future this may be a configurable feature.

function M.clear_completion()
  if current_extmark_id then
    local bufnr = vim.api.nvim_get_current_buf()
    local ns_id = vim.api.nvim_create_namespace('latios')
    vim.api.nvim_buf_del_extmark(bufnr, ns_id, current_extmark_id)
    current_extmark_id = nil
  end
end

This will simply delete the extended mark associated with the completion. current_extmark_id is a global variable the module uses to track the displayed extended marks.

function M.accept_completion()
  local bufnr = vim.api.nvim_get_current_buf()
  local ns_id = vim.api.nvim_create_namespace('latios')
 
  if current_extmark_id then
    local extmark = vim.api.nvim_buf_get_extmark_by_id(bufnr, ns_id, current_extmark_id, { details = true })
    if extmark and extmark[3] then
      local line, col = unpack(extmark)
      local completion_lines = {}
 
      -- Get the first line from virt_text
      if extmark[3].virt_text then
        table.insert(completion_lines, extmark[3].virt_text[1][1])
      end
 
      -- Get additional lines from virt_lines
      if extmark[3].virt_lines then
        for _, virt_line in ipairs(extmark[3].virt_lines) do
          table.insert(completion_lines, virt_line[1][1])
        end
      end
 
      -- Insert the completion text
      local current_line_text = vim.api.nvim_buf_get_lines(bufnr, line, line + 1, false)[1]
      local prefix = string.sub(current_line_text, 1, col)
      local suffix = string.sub(current_line_text, col + 1)
 
      local new_lines = { prefix .. completion_lines[1] }
      for i = 2, #completion_lines do
        table.insert(new_lines, completion_lines[i])
      end
      new_lines[#new_lines] = new_lines[#new_lines] .. suffix
 
      -- Replace the current line and add new lines if necessary
      vim.api.nvim_buf_set_lines(bufnr, line, line + 1, false, new_lines)
 
      -- Move the cursor to the end of the inserted text
      local end_line = line + #new_lines - 1
      local end_col = #new_lines[#new_lines] - #suffix
 
      vim.api.nvim_win_set_cursor(0, { end_line + 1, end_col })
 
      vim.cmd('startinsert!')
    end
    M.clear_completion()
  end
end

This function will take the extended mark text and insert it into the buffer. It will replace the current line, append any additional lines, and move the cursor to the end of the completion. Finally, it’ll remove the extended mark

Challenges and Gotchas

While making this initial MVP there were definitely a couple of challenges I faced along the way. Many of these were related to performance.

Debouncing Completion Requests

As you probably guessed from the fact that a completion is triggered everytime the cursor moves, there are a lot of requests that can occur. That’s why the completion request to server.lua is debounced. There is a short delay before the request actually starts to account for the fact that the cursor will move often while in insert mode. Each time the completion method is called the delay will restart.

Developer Experience

Developing the plugin itself was a bit tricky and cumbersome. I loaded the local plugin using lazy.nvim, but everytime I made a change I had to exit neovim and re-enter the buffer. I will probably invest some more time in just improving the workflow there.

Debugging and Profiling

Similarly, it was hard to debug and profile the plugin. It wasn’t very clear how to print error messages and do a normal print debugging strategy. A lot of the sources I found just recommended printing to a buffer or file, however I ran into some issues while doing that of the copilot being unable to parse the treesitter grammar in the buffer or read it’s filetype. I quickly fixed those, but the workflow definitely still needs some work.

Profiling was also tricky while I was trying to investigate some of the performance issues. To address both of these issues I made a naive debug.lua and profiler.lua modules that probably need some work.

Sources

Next Steps

After reflecting on my development up to this point and writing out my process I can already see a lot of places for improvement. In a lot of my modules I’m calling require at the top instead of locally in functions, while pulling the docs for jobstart I saw it recommended using vim.system instead, a lot of my context.lua probably doesn’t have the most relevant information, and I can probably improve the LLM’s prompt by switching the last message to be a assistant message with the current code line so that the LLM attempts to complete it.

I know a lot of work has also been done in this space in the time I’ve been working on it and I want to keep researching the best methods and approaches to see what I can learn.

  • Research more methods for building a more streamlined context for the copilot
  • Finding better debugging and testing workflows for neovim plugin development
  • Looking into a chat window and a select to edit feature
  • Improve the performance of the plugin

Graph View

Backlinks