はじめに

今日のLLMを扱った開発は、Pythonがそれなりに優遇されている気がします。
OpenAIのライブラリもあり、LangChainのライブラリも初期の頃から提供されていて、僕もそれを利用させてもらっていました。
そんな中で、TypeScriptにのみ提供されているLLMのフレームワークが出現したことを知り、とりあえず見よう見まねで触ってみることにしました。それがMastraというエージェントフレームワークです。
今回は、Mastraを使って簡単なRAGを実装してみます。

Mastraとは

Mastraの公式ページに遷移すると記載がありますが、最近ホットになっているAIエージェントの構築をはじめとして、RAGの構築やLLMOpsなどにも幅広く対応しているフレームワークです。
LLMの開発を初期から支えてきたライブラリとして、LangChainは非常に有名だと思います。僕もLangChainを利用していますが、派生ライブラリの多さは少し面倒でした。

- langchain           -> RAG
- langgraph           -> エージェント
- langchain-community -> サードパーティ製のライブラリを操作
- langsmith, langfuse -> LLMOps

LangChainが初期の頃は、いろんな機能がlangchainにまとまってしまい、それが肥大化し複雑化して来たことで派生させたのだろうとは思います。かと言って、派生ライブラリが増えすぎて、今ではバージョン管理が若干大変です。
そんな中で出現したMastraは、LangChainが派生させた多くの機能を、1つのライブラリで担ってくれるみたいです。これで使い勝手が良好であれば、大変嬉しいです。

まさか、派生しないよね..？

また、LLMの多くは、研究よりも開発を目的として使用されるため、フロントエンド開発でモダンなTypeScriptとの相性を鑑みて、話題になっているみたいです。

僕は普段Pythonしか書かないため、なんでTypeScriptだけなんだ！ と一瞬思ったりしました。ですが、言い換えると、なんでこれまでPythonはLLMにおいてこんなに優遇されてたのだろう？ と思い始めました。Pythonが機械学習に強くなれたのは、先人たちがたまたま機械学習系のライブラリやフレームワークをPython用に作ってくれたからであり、LLM開発もその名残なのでは？と思っています。
OpenAIのライブラリも、当初はPythonとJavaScriptのみの提供でしたが、昨年の冬に.NET / Java / Goの3つが加わったように、これからのLLM開発はPython以外も視野に入れてみるようにします。

環境構築

今回、TypeScriptを使って開発するにあたり、実行環境はDenoを使用しました。
理由としては、普段LLMの開発をする中でJupyter Notebookの形式を愛用しており、似た環境で作業したかったためです。すなわち、REPL形式による開発を実施したく、Denoを選択しました。
ここ1年くらいはZedエディターを使ったコーディングをしていて、TypeScriptの実行環境だとDenoのみREPLに対応しており、Node.jsやBunにはまだ対応してないみたいです。

Ollama環境構築

$ curl -fsSL https://ollama.com/install.sh | sh
$ ollama pull nomic-embed-text

ベクトル化に使用するモデルを、あらかじめダウンロードします。

Deno環境構築

$ brew install deno
$ deno jupyter --unstable

上記を実行することで、Denoが使用できるカーネルを作成できました。

ChromaDB環境構築（起動まで）

公式ドキュメントより

$ pip install chromadb
$ chroma run

ローカル環境にインストールする際には、どうやらpip経由らしいです。pipに限った話ではないですが、詰まるところ、Pythonのパッケージ管理ツールを使用するみたいです。TypeScriptとは一切関係ないみたいです。

このときですが、Python==3.8.10だとエラーになり、ChromaDBが起動できませんでした。どうやらPythonバージョンが古いことが原因だったみたいなので、Python==3.12.3にすると起動できました（古すぎなければ起動すると思います）。

RAGを実装してみる

公式ドキュメント

個人的なあるあるなんですが、LLMサービスは従量課金なので失敗するたびに課金されることに抵抗があります。
そのため可能な場合は、まず最初にローカルLLMを使った実装から始めることが多いです。今回もまずはOllamaを使ってRAGを実装していきます。
その後にOpenAIを使った実装もやっています。

Mastra x Ollama

下準備

example.ts

// %% インポート
import { ChromaVector } from '@mastra/chroma';
import { Document, Metadata } from "@llamaindex/core/schema";
import { EmbeddingModelV1 } from "@ai-sdk/provider";
import { EmbedResult, QueryResult } from '@mastra/core';
import { MDocument } from "@mastra/rag";
import { embed, embedMany } from "ai";
import { ollama } from "ollama-ai-provider";

// ベクトル化モデルを定義
const MODEL_OLLAMA_EMBEDDING: EmbeddingModelV1<string> = ollama.embedding("nomic-embed-text");

example.ts

// %% 検索したいドキュメントを用意（とりあえず過去の自分の記事）
const htmlSource = `<ul data-sourcepos="15:1-22:0">
<li data-sourcepos="15:1-16:123">1段目「<strong>関連</strong>」 → サウナ中に計測できたデータを使って予測する
<ul data-sourcepos="16:5-16:123">
<li data-sourcepos="16:5-16:123">心拍数、体温、サウナ室の温湿度、サウナ利用時間を使って、ととのい具合を予測する</li>
</ul>
</li>
<li data-sourcepos="17:1-18:160">2段目「<strong>介入</strong>」 → サウナ中に計測できたデータの一部を変更する（=条件付けする）ことで、その他がどういう影響を受けるかを予測する
<ul data-sourcepos="18:5-18:160">
<li data-sourcepos="18:5-18:160">あとサウナ室に2分多く入っていると、他のデータ（心拍数、体温など）はどうなるか、ととのい具合はどうなるか</li>
</ul>
</li>
<li data-sourcepos="19:1-22:0">3段目「<strong>反事実</strong>」 → 予想することすらできない状態（=妄想）
<ul data-sourcepos="20:5-22:0">
<li data-sourcepos="20:5-20:170">OOっていうサウナでととのった！XXっていうサウナだったら、さらにととのったのかな？（同時に2つの事象を満たせない）</li>
<li data-sourcepos="21:5-22:0">サウナ室で突然倒れてしまった、そういえば水を全然飲んでなかったので、もし水を飲んでいたらどうなっただろう？（再現性がない）</li>
</ul>
</li>
</ul>`;
const doc: MDocument = MDocument.fromHTML(htmlSource);

今回はRAGの検索対象のドキュメントとして、過去の自分のQiita記事の一部を持って来ています（権利問題無問題）。

RAGの設定

example.ts

// %% チャンク分割を設定
const chunks: Document<Metadata>[] = await doc.chunk({
  strategy: "recursive",
  size: 64,
  overlap: 16
});

今回は文章量がそんなに多くないので、チャンクサイズは小さめにしています。あと、チャンクサイズが大きすぎると検索が大雑把になってしまうので。

example.ts

// %% ベクトル化を設定
const { embeddings } = await embedMany({
  model: MODEL_OLLAMA_EMBEDDING,
  values: chunks.map(chunk => chunk.text)
});

example.ts

// %% ベクトルデータベースを作成
const vectorStore = new ChromaVector({
  path: "http://localhost:8000"
});
await vectorStore.deleteIndex("embeddings");  // 重複エラーを避けるため
await vectorStore.createIndex({
  indexName: "embeddings",
  dimension: 768,
  metric: "cosine"
});
await vectorStore.upsert({
  indexName: "embeddings",
  vectors: embeddings,
  metadata: chunks.map(chunk => ({ id: "サウナと因果推論" })),
  documents: chunks.map(chunk => chunk.text)
});

ChromaDBがデフォルトで8000番ポートにて起動しているので、それに合わせてURLを指定しています。
また、何回かやり直すときに、同一のインデックス名があるとエラーになるため、わざとインデックス名を都度削除しています。ここを削除しないことで、ベクトルデータの永続化が可能です。
また後で触れますが、ベクトル検索時にコサイン類似度を指定しましたが、想像していた指標とは異なるスコアになりました。後々..

ベクトルデータベースを作成すると、UUIDのような文字列が列挙されました。これら1つ1つがベクトル情報になります。

ベクトル検索を実行する

example.ts

// %% ユーザープロンプトをベクトル化
const userQuery = "因果のはしごの3段目は何ですか？"
const queryVector: EmbedResult<string> = await embed({
  model: MODEL_OLLAMA_EMBEDDING,
  value: userQuery
});

example.ts

// %% ベクトル検索を実行
const results: QueryResult[] = await vectorStore.query({
  indexName: "embeddings",
  queryVector: queryVector.embedding,
  topK: 1
});
results

topKというパラメータにて、検索する上位件数を指定でき、今回は1つだけにしています。
結果として、コサイン類似度のスコアは0.375、得られたテキスト情報は→ サウナ中に計測できたデータを使って予測するでした。欲しい情報としては不適切な結果になりました。

Mastra x OpenAI

Ollamaを使って実行確認ができたので、続いてOpenAIを利用してみます。
基本的に流れは全く同じで、はじめの下準備の段階で定義したモデルを変更します。
また、RAGの設定にて指定したdimensionを756→1536に変更します。この2点だけです（APIキーの設定も忘れずに）。

変更点.ts

1. モデルを変更
const MODEL_OLLAMA_EMBEDDING: EmbeddingModelV1<string> = ollama.embedding("nomic-embed-text");

// OpenAIのモデルに変更
const MODEL_OPENAI_EMBEDDING: EmbeddingModelV1<string> = openai.embedding("text-embedding-3-small");

2. ベクトルの次元数を変更
await vectorStore.createIndex({
  indexName: "embeddings",
  dimension: 768,
  metric: "cosine"
});

// 1536次元に変更
await vectorStore.createIndex({
  indexName: "embeddings",
  dimension: 1536,
  metric: "cosine"
});

OpenAIのEmbeddingモデルに変更することで、ローカルに比べると精度がかなり上がると思います。
結果は、コサイン類似度のスコアが0.555、得られたテキスト情報は3段目「反事実」 →となり、無事に欲しい結果を得ることができました！

何かおかしいコサイン類似度

先ほどまでのベクトル検索は上位1件のみを取得していましたが、OpenAIのままで件数を3件に増やしてみた結果が以下です。

あれ？
上位3件を取得してきたはずですが、なぜか先ほどの上位1件の結果よりもコサイン類似度のスコアが高いテキストがある？？
コサイン類似度についてですが一応確認すると、-1 <= cos <= 1の範囲の値を取ることが知られていて、以下のような意味を持っています。

1に近いほど、類似性が高い（=正の相関がある）
-1に近いほど、類似性が真逆（=負の相関がある）
0に近いほど、類似性が皆無（=相関がない）

今回の場合は、1に近いほど情報が類似しているということになるはずですが、なぜか1に近くないテキストが上位1件のときに取得されたみたいです。おかしいですね。

ということで少し調べてみると、今回使用したChromaDBの公式ドキュメントでこのような情報を見つけました。

というわけで、コサイン類似度の計算結果を1から引いた最終的な値を、コサイン類似度としているみたいです。
そのため、先ほどの-1 <= cos <= 1の範囲は0 <= cos_sim <= 2ということになり、意味合いもずれます。

0に近いほど、類似性が高い
1に近いほど、類似性が皆無
2に近いほど、類似性が真逆

これを把握した上で、先ほどの上位3件の件に戻ると、確かに上位1件のコサイン類似度のスコアが低かった点も合点がいきますね。

コサイン類似度による評価指標を正の値（厳密には0は正の値ではないが）に正規化することで、0を目指せば良いのか！と直感的に理解できる方もいるかもしれませんが、個人的にはコサイン類似度はコサイン類似度のままにして欲しいですね..。それとも正規化によって、LLM開発で何かメリットがあるのでしょうか？知見のある方がいらっしゃったら、ぜひ教えて欲しいですmm

おわりに

OpenAIでの動作もでき、簡単にですがMastraを使ってみることができました。RAGというものの、ベクトル検索までしか実施していない点はすみません..。得られた回答をあとはLLMに投げるだけで回答生成までできると思います。
冒頭でも述べましたが、今後のLLM開発はPythonに拘泥することなく、幅広く情報収集したいと思います