Livebook で Ollama と TiDB による RAG チャットを実装する

はじめに

以前の記事で紹介した TiDB 、 Ollama 、 Ruri を組み合わせることで、簡単に日本語 RAG チャットが実装できます

テキスト埋め込みには日本語に特化した Ruri

テキスト生成には比較的軽量で CPU 環境でもそれなりに速く動いてくれる Phi-4 を使用します

実装したノートブックはこちら

TiDB の準備

TiDB Cloud Serverless でクラスターを作成しておきます

本記事の範囲内であれば無料で利用可能です

クラスターが作成されると、以下のような概要画面が表示されます

右上の "Connect" ボタンをクリックしてください

以下のようにデータベースの接続情報が表示されます

右中央の "Generate Password" ボタンをクリックすると、ランダムなパスワードが生成されます

Livebook からはこの接続情報を使って接続します

パスワードはモーダルを閉じると次回以降表示されないので安全な場所に記録しておきましょう

コンテナの準備

Livebook をコンテナで起動し、 Ollama をホスト上で実行する場合、

Ollama はホスト名 host.docker.internal でアクセス可能です

本記事では Livebook と Ollam をコンテナで起動します

docker-compose.with-ollama.yml

---

services:
  livebook_with_ollama:
    image: ghcr.io/livebook-dev/livebook:0.14.7
    container_name: livebook_with_ollama
    ports:
      - '8080:8080'
      - '8081:8081'
    volumes:
      - ./tmp:/tmp

  ollama:
    image: ollama/ollama:0.5.11
    container_name: ollama_for_livebook
    ports:
      - '11434:11434'
    volumes:
      - ollama:/root/.ollama

volumes:
  ollama:

コンテナのビルド・起動

以下のコマンドでコンテナをビルド、起動します

docker compose --file docker-compose.with-ollama.yml up

両方のビルドが完了したら Livebook の URL にアクセスします

Livebook 上での RAG チャットの実装

セットアップ

Livebook で新しいノートブックを開き、セットアップセルで以下のコードを実行します

Mix.install([
  {:myxql, "~> 0.7"},
  {:kino_db, "~> 0.3"},
  {:ollama, "~> 0.8"}
])

今回は Ecto ではなく、簡単に DB 接続、クリ離実行ができる KinoDB を利用します

TiDB は MySQL 互換なので MyXQL で接続します

DB 接続

スマートセルで "Database connection" を選択します

追加されたフォームに TiDB の接続情報を入力してください

セルを実行すると、 TiDB に接続できます

テーブルの作成

スマートセルで "SQL query" を選択します

SQL 実行用のフォームが表示されるので、以下のSQLを入力して実行します

DROP TABLE IF EXISTS rag_index

繰り返し実行する場合のため、前回作ったテーブルを削除するSQLです

もう一つ SQL query のスマートセルを追加してください

時間がかかる場合があるので、 TIMEOUT を 180 秒に設定します

以下の SQL を実行し、 RAG に使用するテーブルを作成します

CREATE TABLE rag_index (
  text VARCHAR(1000),
  embedding VECTOR(768),
  VECTOR INDEX idx_embedding ((VEC_COSINE_DISTANCE(embedding)))
)

text 列に元の文章、 embedding 列に埋め込み表現を格納するようになっています

embedding のサイズは利用するテキスト埋め込みモデルに応じて変更してください

今回使う Ruri-base では 768 になります

チャットの実装

コンテナで起動している Ollama にアクセスするためのクライアントを用意します

URLでしていしている「ollama」の部分は docker-compose で指定したサービス名です

応答に時間がかかることを想定して、 receive_timeout に 300 秒を指定します

client = Ollama.init(base_url: "http://ollama:11434/api", receive_timeout: 300_000)

クライアントで Phi-4 のモデルをダウンロードします

Ollama.pull_model(client, name: "phi4")

読み込みに時間がかかるので、実行前に読み込んでおきます

Ollama.preload(client, model: "phi4")

回答用の関数を用意します

ストリーミングで処理した後、全部のトークンを結合して返すようにしています

answer = fn input, frame ->
  {:ok, stream} =
    Ollama.completion(
      client,
      model: "phi4",
      prompt: input,
      stream: true
    )

  stream
  |> Stream.transform("AI: ", fn chunk, acc ->
    response = acc <> chunk["response"]

    markdown = Kino.Markdown.new(response)
    Kino.Frame.render(frame, markdown)

    {[chunk["response"]], response}
  end)
  |> Enum.join()
end

回答表示用のフレームを用意します

answer_frame = Kino.Frame.new()

大分県の郷土菓子である「やせうま」について質問してみましょう

answer.("やせうまについて100文字程度で説明して", answer_frame)

実行すると、用意しておいたフレームに結果がストリーミングで表示されていきます

しかし、回答内容は典型的なハルシネーションで、そっれぽい嘘です

Phi-4 は「やせうま」を知らないようなので、 RAG で回答できるようにしてあげましょう

テキスト埋め込み

テキスト埋め込み用の Ruri をダウンロードしておきます

Ollama.pull_model(client, name: "kun432/cl-nagoya-ruri-base")

RAG 用の知識として、「やせうま」の情報と、適当な他の文章を用意しておきます

string_inputs = [
  "浦島太郎は日本の昔話の主人公で、亀を助けた礼として竜宮城に招かれます。帰郷時に渡された玉手箱を開けると老人になってしまう物語です。",
  "豆腐小僧は江戸時代の草双紙や錦絵に登場する妖怪で、笠をかぶり盆に乗せた豆腐を持つ子供の姿をしています。特に悪さをせず、愛嬌のある存在として描かれています。",
  "やせうまは大分県の郷土菓子で、茹でた小麦粉の生地にきな粉と砂糖をまぶしたものです。素朴な甘さともちもちした食感が特徴です。"
]

テキスト埋め込み用の関数を用意します

embed = fn input ->
  client
  |> Ollama.embed(
    model: "kun432/cl-nagoya-ruri-base",
    input: input
  )
  |> elem(1)
  |> Map.get("embeddings")
  |> hd()
end

RAG 用の知識をベクトル（埋め込み表現）に変換します

Ruri では文頭に「文章: 」をつけることになっています

embeddings = Enum.map(string_inputs, fn input -> embed.("文章: #{input}") end)

実行結果

[
  [0.0012877434, -0.023524435, 0.041143656, 0.025922664, -0.0028315857, -0.0035114677, -0.041583583,
   -0.029240241, 0.029034998, 0.03579283, -0.012323695, 0.023974337, -0.060204238, -0.044069838,
   0.022237863, -0.0011695586, 0.038902238, 0.0150097655, -0.013371013, 0.03407516, -0.03128297,
   -0.045207165, 0.041558128, -0.02906675, 0.007098051, 0.036273085, -0.027635474, 0.0053589973,
   -0.04444531, -0.045125935, -0.02492561, -0.00773266, 0.019043665, -0.3389603, -0.0029868651,
   -0.036341418, -0.0035760878, -0.06136555, -0.019327547, 0.041891832, 0.05151429, -0.02402123,
   -0.057392303, 0.03515303, -0.0276297, 0.03400351, -0.028113676, 0.034500767, -0.029931104, ...],
  [0.011153108, -0.0013057652, 0.059159093, 0.026615385, 0.0016438881, -0.021420551, -0.043166384,
   -0.015790908, 0.050318047, 0.027289387, -0.015138047, 0.01944528, -3.9178808e-4, -0.026380489,
   0.01622893, 0.01703049, 0.032122057, 0.008805354, -0.10637413, 0.057963602, -0.053950295,
   -0.04365103, 0.034218304, -0.028890342, 0.0031284718, 0.043336585, -0.047029253, -0.011202821,
   -0.0448323, -0.030884938, -0.013659151, 8.4184745e-4, 0.00644853, -0.32884195, -0.025538044,
   -0.016359424, 0.0070622335, -0.020795316, -0.023583861, 0.029994247, 0.035335187, -0.00963662,
   -0.027527787, 0.019642634, -0.046166617, 0.055210374, -0.011101492, 0.031466685, ...],
  [0.027032265, 0.028351044, 0.06167468, 0.040110167, -0.03740636, -0.026319874, -0.025581682,
   -0.012094023, 0.022706969, -0.0071218465, -0.015837645, 0.027824776, -0.037639976, -0.0360941,
   0.019610846, 0.006883774, 0.039032932, 0.04753021, -0.0882825, 0.033915758, -0.052438173,
   -0.066123694, 0.04506811, -0.034176506, 0.009999059, 0.026821274, -0.056069154, 0.0075303135,
   -0.015185168, -0.042341232, -0.018102499, -0.0036201363, 0.014060944, -0.33715412, -0.00935216,
   -0.025623156, 0.005602004, -0.028600829, -0.01685087, 0.047522876, 0.0065040984, -0.008226237,
   -0.01167611, 0.03806767, 0.0057777157, 0.034388494, -0.011633973, ...]
]

ベクトルを RAG 用テーブルに挿入します

[string_inputs, embeddings]
|> Enum.zip()
|> Enum.map(fn {text, embedding} ->
  MyXQL.query!(
    conn,
    ~S"""
    INSERT INTO rag_index (text, embedding) VALUES (?, ?)
    """,
    [text, embedding]
  )
end)

SQL query スマートセルで SELECT * FROM rag_index を実行し、RAG 用テーブルのデータを確認します

テキスト検索

TiDB 上の RAG 用テーブルからテキスト検索する関数を用意します

テキスト検索時は embedding 列とクエリの埋め込み表現のコサイン距離が最も小さい text を取得しています

search = fn query ->
  embedding = embed.("クエリ: #{query}")

  MyXQL.query!(
    conn,
    ~S"""
    SELECT text FROM rag_index
    ORDER BY VEC_COSINE_DISTANCE(embedding, ?)
    LIMIT 1
    """,
    [embedding]
  )
end

「やせうま」についての質問を投げてみます

search.("やせうまについて100文字程度で説明して")

実行結果

ちゃんと「やせうま」の情報が取得できました

RAG チャットの実装

質問を元に TiDB から情報を検索し、それに基づいて返答する関数を用意します

rag = fn input, frame ->
  context =
    input
    |> search.()
    |> Map.get(:rows)
    |> hd()
    |> hd()

  answer.("context: #{context}\n\ncontext に基づいて質問に答えてください\n\n#{input}", frame)
end

入出力用のフォームを用意します

出力途中の回答を表示する stream_frame と、回答がすべて出力された後に表示する output_frame をそれぞれ用意しているのがポイントです

# 出力用フレーム
output_frame = Kino.Frame.new()

# ストリーミング用フレーム
stream_frame = Kino.Frame.new()

# 入力用フォーム
input_form =
  Kino.Control.form(
    [
      input_text: Kino.Input.textarea("メッセージ")
    ],
    submit: "送信"
  )

Kino.Frame.render(output_frame, Kino.Markdown.new(""))
Kino.Frame.render(stream_frame, Kino.Markdown.new(""))

# フォーム送信時の処理
Kino.listen(input_form, fn %{data: %{input_text: input}} ->
  Kino.Frame.append(output_frame, Kino.Markdown.new("あなた: " <> input))
  full_response = rag.(input, stream_frame)
  Kino.Frame.render(stream_frame, Kino.Markdown.new(""))
  Kino.Frame.append(output_frame, Kino.Markdown.new("AI: " <> full_response))
end)

# 入出力を並べて表示
Kino.Layout.grid([output_frame, stream_frame, input_form], columns: 1)

実行結果

ちゃんと質問に答えることができました

まとめ

Ollama と TiDB の組み合わせにより、 Livebook 上で簡単に RAG チャットが実装できました

Livebook ならこれをそのままアプリとして公開することも可能です

Phoenix アプリケーションに組み込む場合は Ecto を使うと良いでしょう