neovim_stream-a-shell-command-into-neovim

Stream a Shell command into Neovim

Introduction

Neovim can read text in using the :r ! command however this doesn’t stream. I needed streaming to integrate with large language models, in my case via Ollama.

A Streamable Command

First set up a command that streams, a simple for i in {1..10}; do echo $i; sleep 1; done will do, but in my case I wanted to stream the output of a python script:

#!/usr/bin/env python3
from ollama import Client
import sys
 
client = Client(host="http://localhost:11434")
 
# message = "Why is the sky blue?"
 
if len(sys.argv) > 1:
    message = sys.argv[1]
 
 
stream = client.chat(
    model="phi3:latest",
    messages=[{"role": "user", "content": message}],
    stream=True,
)
 
for chunk in stream:
    print(chunk["message"]["content"], end="", flush=True)

Streaming it

I figured this out by prompting Codestral (by Mistral.AI) and then reading the documentation for vim.loop.spawn, then trial and error:

First Attempt

This approach does not handle new lines

--- Stream shell command output to buffer
---@param cmd string: The Command to run, e.g. "cat" or "python"
---@param args string[]: The arguments to pass to the command, e.g. "file.py"
local function stream_shell_command(cmd, args, verbose)
  if verbose == nil then
    verbose = false
  end
  -- Pause tracking undos
  -- TODO put in function
  vim.o.undofile = false
  local uv = vim.loop
 
  local stdin = uv.new_pipe()
  local stdout = uv.new_pipe()
  local stderr = uv.new_pipe()
 
  if verbose then
    print("stdin", stdin)
    print("stdout", stdout)
    print("stderr", stderr)
  end
 
 
  local handle, pid = uv.spawn(cmd, {
    args = args, stdio = { stdin, stdout, stderr }
  }, function(code, signal) -- on exit
    print("exit code", code)
    print("exit signal", signal)
  end)
 
  print("process opened", handle, pid)
 
  uv.read_start(stdout, function(err, data)
    assert(not err, err)
    ---@type string data
    if data then
      -- NOTE insert the text with this command
      vim.api.nvim_put({ data }, "c", true, true)
    end
  end)
 
  uv.read_start(stderr, function(err, data)
    assert(not err, err)
    if data then
      print("stderr chunk", stderr, data)
    else
      print("stderr end", stderr)
    end
  end)
 
  uv.write(stdin, "Hello World")
 
  uv.shutdown(stdin, function()
    print("stdin shutdown", stdin)
    -- TODO how to handle nil?
    uv.close(handle, function()
      print("process closed", handle, pid)
    end)
  end)
 
  -- fn_1: https://www.reddit.com/r/neovim/comments/ulx17m/using_nvim_buf_set_text_without_running_into/
end

Fixing new lines

Problem

Unfortunately, this command:

-- NOTE insert the text with this command
vim.api.nvim_put({ data }, "c", true, true)

Does not insert new lines, but using:

vim.api.nvim_put({ data }, "c", true, true)

Introduces too many.

Fixed Function

To fix this I had to look for new lines:

---This function takes in text and puts it into the buffer after the cursor
---
---Unlike vim.api.nvim_put({text}, "c", ... )
---    This handles new lines
---Unlike vim.api.nvim_put({text}, "c", ... )
---    This does not leave a trailing new line in the buffer
---
---This is useful for streaming text that may have new lines
---@param data string: The input string, represents a chunk e.g. `def main():    print(foo)`
local function put_text_with_new_lines(data)
  if data.find(data, "\n") ~= nil then
    -- Split out the new lines
    local text_lines = vim.split(data, "\n")
    for i, text in ipairs(text_lines) do
      vim.schedule(function()
        vim.api.nvim_put({ text }, "c", true, true)
        if i ~= #text_lines then
          vim.api.nvim_put({ "" }, "l", true, false)
        end
      end)
    end
  else
    vim.schedule(function()
      vim.api.nvim_put({ data }, "c", true, true)
    end)
  end
end
 

All together

Putting that all together, the following will stream the output of a shell command into the buffer:

---This function takes in text and puts it into the buffer after the cursor
---
---Unlike vim.api.nvim_put({text}, "c", ... )
---    This handles new lines
---Unlike vim.api.nvim_put({text}, "c", ... )
---    This does not leave a trailing new line in the buffer
---
---This is useful for streaming text that may have new lines
---@param data string: The input string, represents a chunk e.g. `def main():    print(foo)`
local function put_text_with_new_lines(data)
  if data.find(data, "\n") ~= nil then
    -- Split out the new lines
    local text_lines = vim.split(data, "\n")
    for i, text in ipairs(text_lines) do
      vim.schedule(function()
        vim.api.nvim_put({ text }, "c", true, true)
        if i ~= #text_lines then
          vim.api.nvim_put({ "" }, "l", true, false)
        end
      end)
    end
  else
    vim.schedule(function()
      vim.api.nvim_put({ data }, "c", true, true)
    end)
  end
end
 
 
--- Stream shell command output to buffer
---@param cmd string: The Command to run, e.g. "cat" or "python"
---@param args string[]: The arguments to pass to the command, e.g. "file.py"
local function stream_shell_command(cmd, args, verbose)
  if verbose == nil then
    verbose = false
  end
  -- Pause tracking undos
  -- TODO put in function
  vim.o.undofile = false
  local uv = vim.loop
 
  local stdin = uv.new_pipe()
  local stdout = uv.new_pipe()
  local stderr = uv.new_pipe()
 
  if verbose then
    print("stdin", stdin)
    print("stdout", stdout)
    print("stderr", stderr)
  end
 
 
  local handle, pid = uv.spawn(cmd, {
    args = args, stdio = { stdin, stdout, stderr }
  }, function(code, signal) -- on exit
    print("exit code", code)
    print("exit signal", signal)
  end)
 
  print("process opened", handle, pid)
 
  uv.read_start(stdout, function(err, data)
    assert(not err, err)
    ---@type string data
    if data then
      put_text_with_new_lines(data)
    end
  end)
 
  uv.read_start(stderr, function(err, data)
    assert(not err, err)
    if data then
      print("stderr chunk", stderr, data)
    else
      print("stderr end", stderr)
    end
  end)
 
  uv.write(stdin, "Hello World")
 
  uv.shutdown(stdin, function()
    print("stdin shutdown", stdin)
    -- TODO how to handle nil?
    uv.close(handle, function()
      print("process closed", handle, pid)
    end)
  end)
 
  -- fn_1: https://www.reddit.com/r/neovim/comments/ulx17m/using_nvim_buf_set_text_without_running_into/
end

Implementing a Simple LLM Streaming Tool

NOTE

This code assumes \n and not CRLF, adjust for Windows

Wrap the following code into a module and set it to a vmap:

--- Returns the visually selected lines and the end location of those lines
---@return string: The visually selected lines
local function get_lines_as_string(start_line, end_line)
  -- TODO handle selections that aren't linewise
  local selected_lines = vim.fn.getline(start_line, end_line)
 
  if type(selected_lines) == "string" then
    selected_lines = { selected_lines }
  end
 
  return table.concat(selected_lines, "\n")
end
 
local function my_copilot_current_line(clobber)
  if clobber == nil then
    clobber = false
  end
 
  local line = vim.api.nvim_get_current_line()
  -- Move down a line
  local current_line = vim.api.nvim_win_get_cursor(0)[1]
  if clobber then
    vim.api.nvim_buf_set_lines(0, current_line, current_line + 1, false, { "" })
  else
    -- TODO handle buffer of limited length
    vim.api.nvim_win_set_cursor(0, { current_line + 1, 0 })
  end
  stream_shell_command("python", { "/tmp/tmp.CvM0e2ZWPI/file.py", line })
end
 
local function my_copilot_current_selection()
  local indices = get_visual_start_end()
  local start_line, end_line = indices[1], indices[2]
  local content = get_lines_as_string(start_line, end_line)
 
  -- Go to end of visual
  -- vim.api.nvim_win_set_cursor(0, { 0, end_line-1 })
  -- Add a new line
  vim.cmd [[normal o]]
  stream_shell_command("python", { "/tmp/tmp.CvM0e2ZWPI/file.py", content })
end
 
 
-- my_copilot_current_line()
my_copilot_current_selection()