eiji.page

AI時代におけるNeovimでのストリームレスポンス処理と表示

#Neovim #AI

最近はAI関係でストリームレスポンスを扱うことが増えてきました。AIのチャットをパラパラと表示する例のアレです。

Neovimも例外ではありません。AI系のプラグインも増えてきましたが、自分でサクッと書ければカスタマイズの幅が広がって便利です。

この記事では、Neovimでストリームレスポンスを扱う方法をサンプルプログラム付きで解説します。

plenary.nvimのcurl

Neovimでcurlを使うなら、plenary.nvimが便利です。

plenary.nvimは、Neovimで使える便利関数を集めたプラグインです。「よく分からないけど依存プラグインだからインストールしている」という方も多いのでは？

今回は「vim.system関数」とplenary.nvimで実装されている「plenary.curl」の両方を扱います。

環境の準備

まずはplenary.nvimをインストールします。

lazy.nvimだと次のように書いてインストールできます。

require("lazy").setup({
  spec = {
    { "nvim-lua/plenary.nvim", lazy = true },
  },
})

デモ用のAPIサーバーを準備

今回はデモ用のAPIサーバーとしてHonoを使ったTypeScriptのコードを用意します。

index.ts
import { Hono } from "hono";
import { streamText } from "hono/streaming";


const app = new Hono();


app.get("/hello", (c) => {
  return streamText(c, async (stream) => {
    await stream.write("Hello ");
    await stream.sleep(1000);
    await stream.write(`Hono!`);
    await stream.sleep(1000);
    await stream.write(`こんにちは\nこんばんは～`);
  });
});


export default app;

Bunであれば次のコマンド1つでインストールも実行もできます。

Terminal window
bun run index.ts

インストールも簡単なのでこの機会に入れてみては？
ビルドも無しで手軽に実行できるため、筆者は最近めっちゃ使ってます。

実行を確認しつつ重要オプションを学ぶ

次のようにcurlを叩きます。

Terminal window
curl --no-buffer http://localhost:3000/hello

--no-bufferが超重要です。テストだったら絶対出てます。ノートなら赤ペンとかで書いてます。

curlは一定量溜め込んでから表示します。今回はリアルタイムに表示したいため--no-bufferをつけました。

ブラウザでhttp://localhost:3000/helloのようにアクセスしても確認できます。

vim.systemで呼んでみる

vim.systemの例を書きました。

local text = ""


vim.system({ "curl", "--no-buffer", "http://localhost:3000/hello", }, {
  stdout = function(err, data)
    if data then
      text = text .. data
      vim.notify(data)
    end
  end,
}, function()
  -- 最終結果の表示
  vim.notify(text)
end)

実行するにはコマンドモードで次のように入力します。

:luafile %

ちょっとずつ通知されましたね。

plenary.curlで呼んでみる

plenaryのcurlを使ってみましょう。実行時にはvim.systemとの違いに注目しましょう。

local curl = require("plenary.curl")


curl.get("http://localhost:3000/hello", {
  raw = { "--no-buffer" },
  stream = function(error, data)
    if data then
      vim.notify(data)
    end
  end,
  callback = function(response)
    vim.notify(response.body)
  end,
})

今度は一気に通知されましたね。curlでは改行ごとにデータを受け取るようです。

不便に思うかもしれませんが、筆者の扱う範囲だとむしろ改行されているデータを受け取るケースが多いので問題なかったです。
たとえば、Vercel AI SDKのData Stream Protocolsでは次のような感じでデータが飛んできます。

0:"This"
0:" is an"
0:"example."
e:{"finishReason":"stop","usage":{"promptTokens":20,"completionTokens":50},"isContinued":false}
d:{"finishReason":"stop","usage":{"promptTokens":20,"completionTokens":50}}

そんなわけでサーバー側に次の実装を追加します。

index.tsに追加する
app.get("/chat", (c) => {
  return streamText(c, async (stream) => {
    await stream.writeln(`0:"こんにちは！"`);
    await stream.sleep(1000);
    await stream.writeln(`0:"何か"`);
    await stream.sleep(1000);
    await stream.writeln(`0:"お手伝いが"`);
    await stream.sleep(1000);
    await stream.writeln(`0:"必要ですか？\\nそれとも"`);
    await stream.sleep(1000);
    await stream.writeln(`0:"ただの挨拶？"`);
  });
});

writelnで改行くっつけています。

追加したらもう一度サーバーを再起動します。

Terminal window
bun run index.ts

向き先を/chatに変えて実行します。

local curl = require("plenary.curl")


curl.get("http://localhost:3000/chat", {
  raw = { "--no-buffer" },
  stream = function(error, data)
    if data then
      vim.notify(data)
    end
  end,
  callback = function(response)
    vim.notify(response.body)
  end,
})

実践的な処理にかなり近づいてきました。

パースしたりバッファに書き込むのはこの記事の最後にやります。

plenary.curlの使い方

plenary.curlについてもう少し触れておきます。

戻り値のパターン

plenary.curlは引数によって戻り値の型が変わります。ひとつずつ見ていきましょう。

callback・streamなしのパターン

単純にレスポンスだけを受け取りたい場合、次のようにcallbackやstreamを書かずに実行します。

callback・streamなし
local curl = require("plenary.curl")


local response = curl.get("http://localhost:3000/hello")
vim.notify(response.body)

戻り値の型は次のようになっています。

---@class response
---@field exit number The shell process exit code
---@field status number The https response status
---@field headers table The https response headers
---@field body string The http response body

いかにもレスポンスって感じですね。同期的に実行するため、少しだけNeovimの動作が止まります。

callback・streamありのパターン

続いて、callbackやstreamを渡すパターンです。非同期で処理させたいときに使います。

callback・streamあり
local curl = require("plenary.curl")


local job = curl.get("http://localhost:3000/hello", {
  callback = function(response)
    vim.notify(response.body)
  end,
})
vim.print(job) -- 最初に表示されるのはこれ

実行するとテーブル（変数job）が表示されます。callbackかstreamを指定すると、戻り値はplenaryのJobになるのです。

「やっぱりこのリクエストやめよう」となったときに次のように実行すればストップできます。

job:shutdown()
callback

コマンド終了時、コールバックはレスポンスを引数として受け取ります。

---@field callback? fun(response: response): nil
stream

streamはレスポンスを行ごとに受け取ります。

---@field stream? fun(error: string, data: string|nil): nil

dataが実際に受け取るデータです。

dry_runのパターン

dry_runを指定すれば、リクエストを投げずに引数だけが返ってきます。開発中に便利です。

local curl = require("plenary.curl")


local response = curl.get("http://localhost:3000/hello", { dry_run = true })
vim.print(response)

実際に表示された結果は次のとおり。

{ "-sSL", "-D", "/tmp/plenary_curl_951641bb.headers", "--compressed", "-X", "GET", "http://localhost:3000/hello" }

メソッドの指定方法

ここまでの例は全部GETでしたが、当然それ以外のメソッドも叩けます。

curl.get
curl.post
curl.put
curl.head
curl.patch
curl.delete

動的に変更したい場合はcurl.requestを使います。

local curl = require("plenary.curl")


local response = curl.request({
  url = "http://localhost:3000/hello",
  method = "get",
})
vim.notify(response.body)

URLはurl、メソッドはmethodに書きます。筆者はこっちのほうが好きです。

bodyの渡し方

bodyの渡し方もサンプル付きで紹介します。

まずはAPIの追加。JSONを受け取って挨拶します。

index.tsに追加する
app.post("/ohayo", async (c) => {
  const {name} = await c.req.json();
  return c.text(`おはよう ${name}`);
});

サンプルを追加したらもう一度サーバーを再起動します。

Terminal window
bun run index.ts

次のようにbodyに渡します。今回はJSON形式で渡すため、vim.json.encodeでエンコードしてから渡しています。

local curl = require("plenary.curl")


local response = curl.request({
  url = "http://localhost:3000/ohayo",
  method = "post",
  body = vim.json.encode({
    name = "えーたん",
  }),
})
vim.notify(response.body)

バッファにパラパラ書き込みたい

バッファにパラパラと書き込みたい場合は、バッファの最後の位置に文字列を挿入すればよいです。実装としては次のような感じ。

---@param buffer number
---@param text string
function set_text_at_last(buffer, text)
  local lines = vim.split(text, "\n")
  vim.api.nvim_buf_set_text(buffer, -1, -1, -1, -1, lines)
end

実践！

実際にカレントバッファにパラパラ書き込んでみましょう。

処理の例
local curl = require("plenary.curl")


curl.get("http://localhost:3000/chat", {
  raw = { "--no-buffer" },
  stream = function(_, data)
    if data then
      -- パース
      local type_id, content_json = data:match("^([^:]+):(.+)$")
      if type_id ~= "0" then
        return
      end


      local ok, content = pcall(vim.json.decode, content_json)
      if not ok then
        return
      end


      -- データの書き込み
      local lines = vim.split(content, "\n")
      vim.schedule(function()
        vim.api.nvim_buf_set_text(0, -1, -1, -1, -1, lines)
      end)
    end
  end,
  callback = function(response)
    vim.notify(response.body)
  end,
})

バッファの書き込みをvim.scheduleでラップしていますね。こうしないと次のようなエラーがでます。

エラー内容
Error executing luv callback:
lua/write_buffer.lua:17: E5560: nvim_buf_set_text must not be called in a lua loop callback
stack traceback:
  [builtin#36]: at 0x0105aa0e62

一部のAPIはこういったコールバックで直接呼び出せないため、一度ループを抜けた後に実行してもらえるようにスケジュールするのです。

リンク集

細かい話は実装やテストを見ましょう。

「え、面倒……。自分でやったほうが楽だ」という方はvim.systemを使いましょう。

以上、Neovimでストリームレスポンスを扱う方法の解説でした。plenary.curlだと引数を文字列ではなくテーブルのkye-valueで渡せるため、読みやすいです。

senpai.nvimというAIプラグインを開発しており、その過程で発生したアウトプットでした。

#Neovim #AI
記事一覧に戻る

広告

関連記事

すべての記事を見る