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

## はじめに

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

https://qiita.com/RyoWakabayashi/private/58102b89567acf1d3ef6

https://qiita.com/RyoWakabayashi/items/463601c2e455d8b0f9f4

https://qiita.com/RyoWakabayashi/items/8b314f83d4b8a4b7f890

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

https://huggingface.co/cl-nagoya/ruri-base

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

https://huggingface.co/microsoft/phi-4

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

https://github.com/RyoWakabayashi/elixir-learning/blob/main/livebooks/tidb/tidb_rag.livemd

## TiDB の準備

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

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

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

![スクリーンショット 2025-02-17 14.51.41.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/45dc0dfc-744b-4b2c-93bc-e48ddf17750b.png)

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

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

![スクリーンショット 2025-02-17 15.04.57.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/507b1fc7-c198-4e97-ab07-0cbb64154559.png)

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

![スクリーンショット 2025-02-17 15.05.44.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/a92b5243-a104-4501-9be8-7998cc39ad7b.png)

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

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

## コンテナの準備

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

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

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

docker-compose.with-ollama.yml

```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:
```

### コンテナのビルド・起動

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

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

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

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

### セットアップ

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

```elixir
Mix.install([
  {:myxql, "~> 0.7"},
  {:kino_db, "~> 0.3"},
  {:ollama, "~> 0.8"}
])
```
今回は Ecto ではなく、簡単に DB 接続、クリ離実行ができる KinoDB を利用します

https://github.com/livebook-dev/kino_db

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

https://github.com/elixir-ecto/myxql

### DB 接続

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

![スクリーンショット 2025-02-22 15.55.31.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/64a13dcc-6117-49b6-baab-ed4310ac261c.png)

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

![スクリーンショット 2025-02-22 15.56.42.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/4e144986-a309-4d0f-8029-29f36edd4c6e.png)

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

### テーブルの作成

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

![スクリーンショット 2025-02-22 16.00.20.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/4a4c990c-236d-4a06-b316-1582a084aa3b.png)

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

```sql
DROP TABLE IF EXISTS rag_index
```

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

![スクリーンショット 2025-02-23 9.13.55.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/776ef8ee-ab9f-4c05-bd38-ab5084ec4752.png)

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

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

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

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

![スクリーンショット 2025-02-23 9.17.43.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/7430b5b7-1f29-4e54-8545-3f48077ca2a9.png)

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

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

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

### チャットの実装

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

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

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

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

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

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

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

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

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

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

```elixir
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
```

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

```elixir
answer_frame = Kino.Frame.new()
```

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

https://www.maff.go.jp/j/keikaku/syokubunka/k_ryouri/search_menu/menu/yaseuma_oita.html

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

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

![yaseuma.gif](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/e9b24f1d-da35-492e-a6eb-582c093f067f.gif)

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

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

### テキスト埋め込み

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

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

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

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

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

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

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

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

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

実行結果

```elixir
[
  [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 用テーブルに挿入します

```elixir
[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 用テーブルのデータを確認します

![スクリーンショット 2025-02-23 10.22.25.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/e8a282c2-a154-4a5a-84c3-b247f5cad671.png)

### テキスト検索

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

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

```elixir
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
```

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

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

実行結果

![スクリーンショット 2025-02-23 10.27.26.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/2bc145d4-fda2-44ce-9c12-a69a026d5d36.png)

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

### RAG チャットの実装

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

```elixir
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` をそれぞれ用意しているのがポイントです

```elixir
# 出力用フレーム
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)
```

実行結果

![yaseuma_rag.gif](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/8d42f985-d1d3-462b-9264-ffb383ebaa5c.gif)

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

## まとめ

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

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

https://qiita.com/RyoWakabayashi/items/8ef32cb58b339291ca5d

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

https://qiita.com/RyoWakabayashi/items/58102b89567acf1d3ef6