🐈

Firebase GenkitでRAGを使う

2025/01/31に公開

この記事では、Firebase Genkitを使ってRAG(Retrieval-Augmented Generation)を実装する方法を簡単に紹介します。

Firebase Genkitとは

Firebase Genkitは、Firebaseが提供するLLM開発の統合環境です。
LLMの実行、プロンプト管理、トレースなど、さまざまなツールをFirebaseや他のツールと連携して利用することが可能です。

良ければこちらの記事も参照ください。

https://zenn.dev/hakoten/articles/dfa8750518fddf

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)
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)
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