Firebase GenkitでRAGを使う
この記事では、Firebase Genkitを使ってRAG(Retrieval-Augmented Generation)を実装する方法を簡単に紹介します。
Firebase Genkitとは
Firebase Genkitは、Firebaseが提供するLLM開発の統合環境です。
LLMの実行、プロンプト管理、トレースなど、さまざまなツールをFirebaseや他のツールと連携して利用することが可能です。
良ければこちらの記事も参照ください。
RAGのために用意されている機能
Firebase GenkitはRAGを簡単に実装できるように、主に以下の3つの機能を提供しています。
機能名 | 役割 |
---|---|
インデクサー (Indexer) | - ベクトルデータベースを使ってドキュメントのインデックスを作成する |
エンベッダー (Embedder) | - ドキュメントやテキストをベクトル表現に変換することで、類似度検索を可能にする - 専用のMLモデルを使用する |
リトリーバー (Retriever) | - エンベッダーによって数値化されたデータをもとに、ユーザーの質問文(クエリ)に対して最適な情報を検索 - 類似度スコアが高いチャンクを上位N件などで返却し、最終的な回答生成(RAG)のベースとなる情報を提供 |
RAGを一通り実装
基本は、公式ページのサンプルに沿って進めます。
今回は、総務省が公開している「AI進展の経緯と生成AIのインパクト」というPDFを使用します。
インデクサー、エンベッダーを使ってベクトルデータベースを作成
ベクトルデータベースには、更新サンプルと同様に「devLocalVectorstore」というローカル環境用の開発専用ストアを使います。
まず、検索対象となるドキュメントをベクトルデータベースに保存する必要があります。
ベクトルデータベースへ保存するためのツール
ベクトルデータベースへ保存するには、以下の手順を実行する必要があります。
※ 2.PDFをテキストに変換には、pdf-parse。
※ 3.PDFをチャンク化には、llm-chunkというオープンソースのライブラリを使います。
全体コード
まずは、全体のコードを以下に示します。
全体のコード(indexer.ts)
import { Document } from "genkit/retriever";
import { chunk } from "llm-chunk";
import { readFile } from "fs/promises";
import path from "path";
import pdf from "pdf-parse";
import { run } from "genkit";
import {
devLocalIndexerRef,
devLocalVectorstore,
} from "@genkit-ai/dev-local-vectorstore";
import { textEmbedding004, vertexAI } from "@genkit-ai/vertexai";
import { z, genkit } from "genkit";
const ai = genkit({
plugins: [
// embedderのモデル「textEmbedding004」を使用するためのVertexAIを設定する
vertexAI(),
// ベクターストア(devLocalVectorstore)の設定、embedderにはtextEmbedding004モデルを指定する
devLocalVectorstore([
{
indexName: "aiDoc",
embedder: textEmbedding004,
},
]),
],
});
// PDFファイルからテキストを抽出する関数
async function extractTextFromPdf(filePath: string) {
const pdfFile = path.resolve(filePath);
const dataBuffer = await readFile(pdfFile);
// pdf-parseを使用してPDFファイルからテキストを抽出する
const data = await pdf(dataBuffer);
return data.text;
}
// ベクターストアのDocumentを保存するためのインデクサー
export const aiDocIndexer = devLocalIndexerRef("aiDoc");
// フローの定義
export const indexAiDoc = ai.defineFlow(
{
name: "indexAiDoc",
// 入力はString
inputSchema: z.string().describe("PDF file path"),
// 出力はvoid
outputSchema: z.void(),
},
async (filePath: string) => {
filePath = path.resolve(filePath);
// PDFファイルからテキストを抽出する
const pdfTxt = await run("extract-text", () =>
extractTextFromPdf(filePath)
);
// テキストをチャンクに分割する
const chunks = await run("chunk-it", async () =>
chunk(pdfTxt, {
minLength: 1000,
maxLength: 2000,
splitter: "sentence",
overlap: 100,
delimiters: "",
})
);
console.log(chunks);
// チャンクをDocumentに変換してベクターストアのDocumentに変換
const documents = chunks.map((text) => {
return Document.fromText(text, { filePath });
});
// ベクターストアに保存する
await ai.index({
indexer: aiDocIndexer,
documents,
});
}
);
コードの解説
const ai = genkit({
plugins: [
// embedderのモデル「textEmbedding004」を使用するためのVertexAIを設定する
vertexAI(),
// ベクターストア(devLocalVectorstore)の設定、embedderにはtextEmbedding004モデルを指定する
devLocalVectorstore([
{
indexName: 'aiDoc',
embedder: textEmbedding004,
},
]),
],
});
Genkitのアプリインスタンスを初期化しています。ベクターストアへ保存するモデルとして「textEmbedding004」を利用しています。
// PDFファイルからテキストを抽出する関数
async function extractTextFromPdf(filePath: string) {
const pdfFile = path.resolve(filePath);
const dataBuffer = await readFile(pdfFile);
// pdf-parseを使用してPDFファイルからテキストを抽出する
const data = await pdf(dataBuffer);
return data.text;
}
これは公式サンプルと同様で、ファイルパスを指定してpdf-parseを使用し、PDFからテキストを抽出する関数です。
// ベクターストアのDocumentを保存するためのインデクサー
export const aiDocIndexer = devLocalIndexerRef("aiDoc");
// フローの定義
export const indexAiDoc = ai.defineFlow(
{
name: "indexAiDoc",
// 入力はString
inputSchema: z.string().describe("PDF file path"),
// 出力はvoid
outputSchema: z.void(),
},
async (filePath: string) => {
filePath = path.resolve(filePath);
// PDFファイルからテキストを抽出する
const pdfTxt = await run("extract-text", () =>
extractTextFromPdf(filePath)
);
console.log(pdfTxt);
// テキストをチャンクに分割する
const chunks = await run("chunk-it", async () =>
chunk(pdfTxt, {
minLength: 1000,
maxLength: 2000,
splitter: "sentence",
overlap: 100,
delimiters: "",
})
);
console.log(chunks);
// チャンクをDocumentに変換してベクターストアのDocumentに変換
const documents = chunks.map((text) => {
return Document.fromText(text, { filePath });
});
// ベクターストアに保存する
await ai.index({
indexer: aiDocIndexer,
documents,
});
}
);
これが、Genkitのフローを利用してベクターストアに保存する処理です。
const pdfTxt = await run("extract-text", () =>
extractTextFromPdf(filePath)
);
run
メソッドはフローステップ(フローステップ)の一つとして処理をトレースできるようにするための仕組みです。run
でステップをラップすることで、GenkitのGenkit Developer UI上で各ステップの入力・出力などを確認可能です。
const chunks = await run("chunk-it", async () =>
chunk(pdfTxt, {
minLength: 1000,
maxLength: 2000,
splitter: "sentence",
overlap: 100,
delimiters: "",
})
);
llm-chunkを使ってPDFから抽出したテキストをチャンク化します。llm-chunk
の設定値で分割の仕方を指定できます。
// ベクターストアのDocumentを保存するためのインデクサー
export const aiDocIndexer = devLocalIndexerRef("aiDoc");
...
// チャンクをDocumentに変換してベクターストアのDocumentに変換
const documents = chunks.map((text) => {
return Document.fromText(text, { filePath });
});
// ベクターストアに保存する
await ai.index({
indexer: aiDocIndexer,
documents,
});
最後に、分割したチャンクをベクターストアに保存できる形(Documentインスタンス)に変換し、インデクサーを指定して保存します。
Document.fromText(text, { filePath })
の filePath は、そのテキストと一緒に保存するメタデータになります。
実行
GenkitのCLIツールから、このフローを実行します。
(Genkit デベロッパー UIの起動)
genkit start -- pnpm tsx indexer.ts
(フローの実行 ※別ターミナルで実行)
genkit flow:run indexAiDoc '"ai.pdf"'
devLocalIndexerRef
のインデクサーを使用すると、ローカルディレクトリに _db<name>.json というベクターデータファイルが作成されます。
リトリーバーを使ってドキュメントからデータを取得
全体コード
まずは、全体のコードを以下に示します。
全体のコード(retriever.ts)
import { textEmbedding004, vertexAI, gemini15Flash } from "@genkit-ai/vertexai";
import { z, genkit } from "genkit";
import {
devLocalVectorstore,
devLocalRetrieverRef,
} from "@genkit-ai/dev-local-vectorstore";
const ai = genkit({
plugins: [
// embedderのモデル「textEmbedding004」リトリーバーで使用するモデル「gemini15Flash」を使用するためのVertexAIを設定する
vertexAI(),
// データを保存したベクターストアを参照するためにdevLocalVectorstoreを設定する
devLocalVectorstore([
{
indexName: "aiDoc",
embedder: textEmbedding004,
},
]),
],
});
// ベクターストアを参照するためのリトリーバーを設定する
export const aiDocRetriever = devLocalRetrieverRef("aiDoc");
// フローの定義
export const aiDocFlow = ai.defineFlow(
{ name: "aiQA", inputSchema: z.string(), outputSchema: z.string() },
async (input: string) => {
// リトリーバーを使ってベクターストアを参照
const docs = await ai.retrieve({
retriever: aiDocRetriever,
query: input,
options: { k: 3 },
});
// リトリーバーからの参照結果をLLMに渡して回答を生成する
const { text } = await ai.generate({
prompt: `
あなたはAIのドキュメントに関する質問に答えることができるサポーターです。
AIアシスタントとして行動しています。
提供されたコンテキストのみを使用して質問に答えてください。
わからない場合は、答えを作り上げないでください。
質問: ${input}`,
docs,
model: gemini15Flash,
});
return text;
}
);
コードの解説
const ai = genkit({
plugins: [
// embedderのモデル「textEmbedding004」リトリーバーで使用するモデル「gemini15Flash」を使用するためのVertexAIを設定する
vertexAI(),
// データを保存したベクターストアを参照するためにdevLocalVectorstoreを設定する
devLocalVectorstore([
{
indexName: "aiDoc",
embedder: textEmbedding004,
},
]),
],
});
インデックス作成時と同様に、Genkitのアプリインスタンスを初期化しています。
// ベクターストアを参照するためのリトリーバーを設定する
export const aiDocRetriever = devLocalRetrieverRef("aiDoc");
// フローの定義
export const aiDocFlow = ai.defineFlow(
{ name: "aiQA", inputSchema: z.string(), outputSchema: z.string() },
async (input: string) => {
// リトリーバーを使ってベクターストアを参照
const docs = await ai.retrieve({
retriever: aiDocRetriever,
query: input,
options: { k: 3 },
});
// リトリーバーからの参照結果をLLMに渡して回答を生成する
const { text } = await ai.generate({
prompt: `
あなたはAIのドキュメントに関する質問に答えることができるサポーターです。
AIアシスタントとして行動しています。
提供されたコンテキストのみを使用して質問に答えてください。
わからない場合は、答えを作り上げないでください。
質問: ${input}`,
docs,
model: gemini15Flash,
});
return text;
}
);
こちらが、RAGの処理を実際に実行するフローです。
// リトリーバーを使ってベクターストアを参照
const docs = await ai.retrieve({
retriever: aiDocRetriever,
query: input,
options: { k: 3 },
});
ローカルのベクターストアから類似度の高いデータを取得しています。k: 3
で上位3件を返します。
// リトリーバーからの参照結果をLLMに渡して回答を生成する
const { text } = await ai.generate({
prompt: `
あなたはAIのドキュメントに関する質問に答えることができるサポーターです。
AIアシスタントとして行動しています。
提供されたコンテキストのみを使用して質問に答えてください。
わからない場合は、答えを作り上げないでください。
質問: ${input}`,
docs,
model: gemini15Flash,
});
リトリーバーから受け取ったdocs
をLLMに渡し、回答を生成します。ここでは「gemini15Flash」というモデルを使用しています。
実行
GenkitのCLIツールから、このフローを実行します。
(Genkit デベロッパー UIの起動)
genkit start -- pnpm tsx retriever.ts
(フローの実行 ※別ターミナルで実行)
genkit flow:run indexAiDoc '"ai.pdf"'
genkit flow:run aiQA '"生成 AI による経済効果について教えて"'
...
Telemetry API running on http://localhost:4034
Running '/flow/aiQA' (stream=false)...
Result:
"生成AIは経済に大きな影響を与える可能性を秘めています。\n\n* 生成AIは、従来AIが適用しづらかった業務領域も含めて、コンテンツ制作、カスタマーサポート、建設分野など、様々な業務領域での 業務の変革を可能にします。\n* 2023年3月17日、OpenAIとペンシルバニア大学が発表した論文によれば、80%の労働者が、彼らの持つタスクのうち少なくとも10%が大規模言語モデルの影響を受け、その うち19%の労働者は、50%のタスクで影響を受ける。\n* ボストンコンサルティンググループの分析によると、生成AIの市場規模は、2027年に1,200億ドル規模になると予想されています。最も大きな市場 は「金融・銀行・保険」で、次に「ヘルスケア」、「コンシューマー」と続きます。"
flow:run aiQA '"料理について教えて"'
...
Telemetry API running on http://localhost:4034
Running '/flow/aiQA' (stream=false)...
Result:
"申し訳ありませんが、このドキュメントには料理に関する情報はありません。"
終わりに
Firebase Genkitは、ベクトルデータベースの構築やリトリーバー機能を始めとする、RAGを実装するための仕組みを簡単にアクセスできるAPIとして提供しています。
ベクターストアとしてFireStoreを標準提供で使えるなど、Firebase Tool郡との親和性も強みです。この記事が誰かの参考になれば幸いです。
Discussion