前回Mastraを使ってAIにしりとりをしてもらったのですが、本当にmastraの意味がなくて備忘録としてあまり機能していなかったと思ったのでちゃんとtoolsを使うRAGにしたいと思います
RAGについて
本記事では類似近傍(今回はコサイン類似度で)を取得し、その結果を用いて生成AIに回答してもらうことをRAGと表現します。
ということで今回は単語一覧を用意して、その中で類似している物をピックアップし、その中からAIに選んでもらうようなしりとりを作っていきます。それって類似近傍だけでAI要らなくね
toolを作る
先にtoolの部分を説明します。
// 類似単語取得ツールの定義
export const getSimilarityWordTool = async () => {
return createTool({
id: "get-similarity-word",
description:
"しりとりの回答として使用できる、入力した単語と類似した単語を上位5件取得します。",
inputSchema: z.object({
word: z.string().describe("Input word"),
}),
outputSchema: z.object({
words: z
.array(
z.object({
word: z.string().describe("Similar word"),
read: z.string().describe("Reading of the similar word"),
})
)
.describe("Similar words"),
}),
execute: async ({ context }) => {
return { words: await getSimilarWords(context.word) };
},
});
};
こんな感じでtoolを作ります。AIが使用する際に実行される部分はexecute関数です。
このツールの場合はgetSimilarWordsを呼び出して、その返り値をAIが使う感じです。
inputSchemaとoutputSchemaも大事なのですが、toolの名前とdescriptionはAIがツールを使う上で判断するのでAIがちゃんとわかりやすい様にしておきましょう。
ベクトルデータベースの作成
今回はDB作るほどではないと判断しjson形式でファイル保存しました。ちゃんとやるならpgvectorやQdrantを使いましょう。(何回もベクトルデータベース作ったらお金いっぱいかかるかな～って思って保存したけど多分そんな事なかった)
/**
* ベクトルデータベース(ファイル)の作成
* @return 作成したベクトルデータベース
*/
const createVectorDatabase = async (
text: string[]
): Promise<Array<Embedding>> => {
console.log("Creating vector database...");
const response = await openai.embeddings.create({
model: "text-embedding-3-small",
input: text,
});
// ベクトルデータベース(ファイル)の作成
fs.writeFile(VECTOR_DB_FILE_PATH, JSON.stringify(response.data), "utf-8");
return response.data;
};
openai.embeddings.createでモデル指定して、ベクトル化したいデータを配列で渡せばベクトル化されたデータが配列で返ってきます。
類似近傍の取得
// 類似単語取得関数
const getSimilarWords = async (
word: string
): Promise<{ word: string; read: string }[]> => {
// ベクトルデータベースの取得
const vectorDb = await getVectorDb();
// 単語データ(カテゴリ)の取得
const wordDataList = await readCsvFile();
// 入力単語のベクトルを取得
const response: CreateEmbeddingResponse = await openai.embeddings.create({
model: "text-embedding-3-small",
input: word,
});
const inputVector = response.data;
// コサイン類似度を計算して類似度を格納
const similarities = vectorDb.map((v, index) => {
const vector = v.embedding; // ベクトルを取得
const similarity = cosineSimilarity(inputVector[0].embedding, vector); // コサイン類似度を計算
return { similarity, category: wordDataList[index] }; // 類似度とカテゴリを格納
});
// 類似度でソートし、上位5件を取得
similarities.sort((a, b) => b.similarity - a.similarity);
const top5 = similarities.slice(0, 5);
console.log(top5);
return top5.map((p) => {
return { word: p.category[0], read: p.category[1] };
});
};
やってることは検索したいデータ(今回の場合は単語)の埋め込みベクトルを取得して、ベクトルデータベース内の物とのコサイン類似度を算出し上位5件を取ってくるって感じです。
(以下のコサイン類似度の計算する関数はAIに作ってもらいました。ありがとうございます)
/**
* コサイン類似度を計算する関数
* @param a ベクトルA
* @param b ベクトルB
* @returns コサイン類似度
*/
const cosineSimilarity = (a: number[], b: number[]): number => {
let dotProduct = 0;
let normA = 0;
let normB = 0;
for (let i = 0; i < a.length; i++) {
dotProduct += a[i] * b[i]; // 内積
normA += a[i] ** 2; // Aのノルム
normB += b[i] ** 2; // Bのノルム
}
normA = Math.sqrt(normA);
normB = Math.sqrt(normB);
// ゼロ除算を防ぐ
if (normA === 0 || normB === 0) return 0;
return dotProduct / (normA * normB); // コサイン類似度
};
toolを動かしてみる
Mastra playgroundかのサイドバーらtoolの単体実行も出来ます。
作ったtoolが一覧に表示されない場合は何かしらのエージェントに設定されているかどうか確認してください
エージェントに設定されていないtoolは一覧に表示されないっぽい？
コンソールの方でも正しく近傍が取れてるみたいです。
エージェントに組み込む
const tools = {
similarityWordTool: await getSimilarityWordTool(),
};
// RAGしりとりエージェント
export const ragWordChainAgent = new Agent({
name: "Word Chain Agent",
instructions: `
あなたはしりとりをするエージェントです。
しりとりで使う単語はsimilarityWordToolを必ず使って取得してください。
similarityWordToolで得られた単語の中からしりとりのルールに従った単語を選んでください。
出力結果は必ずJSON形式で以下のフォーマットに従ってください。
{
"word": "次の単語",
"status": "in the game" // "in the game" または "game over"
}
`,
tools: tools,
model: openai("gpt-5-nano"),
});
前回と違う箇所はinstructionsのシステムプロンプトでsimilarityWordToolを使うように指示を出しているのとtoolsを設定している箇所だけです。toolsの設定は忘れずに
実行してみる
「りんご」で始めて、すぐ終わったと思ったら・・・
類似している5件の中に「ご」で始まる言葉がなくて、何も返せずに終わってますね。
"similarityWordToolから単語を選ぶ"という指示を守っていますね。RAGの特徴が出ているとも言えます。
これについての修正ですが、面倒だったので類似しているTOP50を取得するに実装を変えました。先頭の文字でフィルターをかけろ
50個もあればしりとり繋がるっしょ
// 類似度でソートし、上位50件を取得
similarities.sort((a, b) => b.similarity - a.similarity);
const top50 = similarities.slice(0, 50);
console.log(top50);
return top50.map((p) => {
return { word: p.category[0], read: p.category[1] };
});
では再度実行してみましょう。先ほどと同じく「りんご」でスタートさせます。
今回は4回目まで続きました。3回目までは正しく類似リストの中から取得していて良いですね。
まあ4回目で言葉が見当たらずに終了してしまいましたがこれは「る攻め」が成功したということで・・・
今回、RAGを使ったしりとりに拡張しようと思って構想を練っている時に「うん、こりゃあ完全に類似度近傍だけでよいでしょう！」という謎の声が聞こえて来ましたが、AIを間に挟む事によって「る攻め」が刺さるなど個人的にはより人間の様なしりとりになったと思っています。
前回と今回作成した物は以下に公開してあります。