検索拡張生成（RAG）

Firebase Genkit は、検索拡張生成（RAG）フローの構築に役立つ抽象化と、関連ツールとの統合を提供するプラグインを提供します。

RAG とは

検索拡張生成は、外部の情報源を LLM のレスポンスに組み込むために使用される手法です。LLM は通常、広範な資料に基づいてトレーニングされますが、LLM の実用には特定分野の知識が必要になることが多いため、これを可能にすることが重要です（たとえば、LLM を使用して、自社の製品に関する顧客からの質問に回答するなど）。

解決策の 1 つは、より具体的なデータを使用してモデルをファインチューニングすることです。ただし、コンピューティングの費用面と、十分なトレーニング データを準備するために必要な労力面の双方で、費用がかかる可能性があります。

一方、RAG は、モデルに渡されるときに外部データソースをプロンプトに組み込むことで機能します。たとえば、「バートとリサの関係は？」というプロンプトが、関連する情報を先頭に追加して拡大（「拡張」）され、「ホーマーとマージの子供はバート、リサ、マギーです。バートとリサの関係は？」というプロンプトになります。

このアプローチにはいくつかのメリットがあります。

• モデルを再トレーニングする必要がないため、費用対効果が高めることができます。
• データソースを継続的に更新でき、LLM は更新された情報をすぐに使用できます。
• LLM の回答で参照を引用できるようになります。

一方、RAG を使用すると、当然プロンプトが長くなり、一部のLLM API サービスでは、送信する入力トークンごとに料金が発生します。最終的には、アプリケーションの費用に関するトレードオフを評価する必要があります。

RAG は非常に広い分野であり、最高品質の RAG を実現するためにさまざまな手法が使用されます。Genkit のコア フレームワークは、RAG を実行するための 3 つの主要な抽象化を提供します。

• インデクサー: ドキュメントを「インデックス」に追加します。
• エンベッダー: ドキュメントをベクトル表現に変換します。
• リトリーバー: クエリを指定して「インデックス」からドキュメントを検索します。

これらの定義には、意図的に曖昧さを残しています。これは、Genkit では「インデックス」とは何か、またはインデックスからドキュメントがどのように正確に検索されるかを、一つに定めないようにしているためです。Genkit は Document 形式のみを提供し、それ以外はすべて、リトリーバーまたはインデクサーの実装プロバイダによって定義されます。

インデクサ

インデックスは、特定のクエリに関連するドキュメントをすばやく検索できるように、ドキュメントを追跡する役割を果たします。ほとんどの場合、これはベクトル データベースを使用して実現されます。ベクトル データベースは、エンベディングと呼ばれる多次元ベクトルを使用してドキュメントのインデックスを作成します。テキスト エンベディングは、テキストの文章で表現される概念を（不透明な形で）表現したものであり、専用の ML モデルを使用して生成されます。ベクトル データベースは、そのエンベディングを使用してテキストのインデックスを作成することで、概念的に関連するテキストをクラスタ化し、新しいテキスト文字列（クエリ）に関連するドキュメントを検索できます。

生成のためにドキュメントを検索するには、ドキュメントをドキュメント インデックスに取り込む必要があります。一般的な取り込みフローは次のとおりです。

1. 関連する部分のみがプロンプトの強化に使用されるように、大きなドキュメントを小さなドキュメントに分割します（チャンク化）。多くの LLM はコンテキスト ウィンドウが限られているため、プロンプトにドキュメント全体を含めることは現実的ではないことが、これを行う理由です。

   Genkit には組み込みのチャンク ライブラリはありませんが、Genkit と互換性のあるオープンソース ライブラリがあります。

2. チャンクごとにエンベディングを生成します。使用しているデータベースに応じて、エンベディング生成モデルで明示的に行うことも、データベースが提供するエンベディング生成ツールを使用することもできます。

3. テキスト チャンクとそのインデックスをデータベースに追加します。

安定したデータソースを使用している場合は、取り込みフローを頻繁に実行しないか、1 回だけ実行することをおすすめします。一方、頻繁に変更されるデータを処理する場合は、取り込みフローを継続的に実行することをおすすめします（たとえば、ドキュメントが更新されるたびに Cloud Firestore トリガーで実行する）。

エンベッダー

エンベッダーは、コンテンツ（テキスト、画像、音声など）を受け取り、元のコンテンツの意味をエンコードする数値ベクトルを作成する関数です。前述のように、エンベッダーはインデックス作成プロセスの一部として活用されますが、インデックスなしでエンベディングを作成するために独立に使用することもできます。

リトリーバー

リトリーバーは、あらゆる種類のドキュメント検索に関連するロジックをカプセル化するコンセプトです。検索の最も一般的なケースには、通常、ベクトル ストアからの検索が含まれますが、Genkit では、データを返す任意の関数をリトリーバーにすることができます。

リトリーバーの作成には、提供されている実装のいずれかを使用するか、独自に作成します。

サポートされているインデクサー、リトリーバー、エンベッダー

Genkit は、プラグイン システムでインデクサーとリトリーバーをサポートします。公式にサポートされているプラグインは次のとおりです。

• Cloud Firestore ベクトル ストア
• Vertex AI Vector Search
• Chroma DB ベクトル データベース
• Pinecone クラウド ベクトル データベース

また、Genkit は、事前定義されたコード テンプレートで次のベクトル ストアをサポートします。テンプレートは、データベースの構成とスキーマに合わせてカスタマイズできます。

• pgvector を使用した PostgreSQL

エンベディング モデルのサポートは、次のプラグインによって提供されます。

プラグイン	モデル
Google Generative AI	Gecko テキスト エンベディング
Google Vertex AI	Gecko テキスト エンベディング

RAG フローを定義する

次の例は、レストラン メニューの PDF ドキュメントのコレクションをベクトル データベースに取り込み、フロー内で使用するために検索する方法を示しています。

PDF の処理に必要な依存関係をインストールする

npm install llm-chunk pdf-parse @genkit-ai/dev-local-vectorstore
npm i -D --save @types/pdf-parse

ローカル ベクター ストアを構成に追加する

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: [
    // vertexAI provides the textEmbedding004 embedder
    vertexAI(),

    // the local vector store requires an embedder to translate from text to vector
    devLocalVectorstore([
      {
        indexName: 'menuQA',
        embedder: textEmbedding004,
      },
    ]),
  ],
});

インデクサーを定義する

次の例は、PDF ドキュメントのコレクションを取り込んでローカルのベクトル データベースに保存するインデクサーの作成方法を示しています。

シンプルなテストとプロトタイピングのために、Genkit が提供するローカルのファイルベースのベクトル類似度リトリーバーを使用します（本番環境では使用しないでください）。

インデクサーを作成する

export const menuPdfIndexer = devLocalIndexerRef('menuQA');

チャンク化構成を作成する

この例では、シンプルなテキスト スプリッターを提供する llm-chunk ライブラリを使用して、ドキュメントをベクトル化可能なセグメントに分割します。

次の定義では、チャンク化関数を構成して、1, 000 ～ 2, 000 文字のドキュメント セグメントを句の末尾で分割し、チャンクのオーバーラップを 100 文字にします。

const chunkingConfig = {
  minLength: 1000,
  maxLength: 2000,
  splitter: 'sentence',
  overlap: 100,
  delimiters: '',
} as any;

このライブラリのその他のチャンク オプションについては、llm-chunk のドキュメントをご覧ください。

インデクサ フローを定義する

import { Document } from 'genkit/retriever';
import { chunk } from 'llm-chunk';
import { readFile } from 'fs/promises';
import path from 'path';
import pdf from 'pdf-parse';

async function extractTextFromPdf(filePath: string) {
  const pdfFile = path.resolve(filePath);
  const dataBuffer = await readFile(pdfFile);
  const data = await pdf(dataBuffer);
  return data.text;
}

export const indexMenu = ai.defineFlow(
  {
    name: 'indexMenu',
    inputSchema: z.string().describe('PDF file path'),
    outputSchema: z.void(),
  },
  async (filePath: string) => {
    filePath = path.resolve(filePath);

    // Read the pdf.
    const pdfTxt = await run('extract-text', () =>
      extractTextFromPdf(filePath)
    );

    // Divide the pdf text into segments.
    const chunks = await run('chunk-it', async () =>
      chunk(pdfTxt, chunkingConfig)
    );

    // Convert chunks of text into documents to store in the index.
    const documents = chunks.map((text) => {
      return Document.fromText(text, { filePath });
    });

    // Add documents to the index.
    await ai.index({
      indexer: menuPdfIndexer,
      documents,
    });
  }
);

インデクサ フローを実行する

genkit flow:run indexMenu "'menu.pdf'"

indexMenu フローを実行すると、ベクトル データベースにドキュメントが供給され、検索ステップを含む Genkit フローでの使用が可能になります。

検索を含むフローを定義する

次の例は、RAG フロー内でリトリーバーを使用する方法を示しています。インデクサーの例と同様に、この例では Genkit のファイルベースのベクトル リトリーバーを使用します。これは本番環境では使用しないでください。

import { devLocalRetrieverRef } from '@genkit-ai/dev-local-vectorstore';

// Define the retriever reference
export const menuRetriever = devLocalRetrieverRef('menuQA');

export const menuQAFlow = ai.defineFlow(
  { name: 'menuQA', inputSchema: z.string(), outputSchema: z.string() },
  async (input: string) => {
    // retrieve relevant documents
    const docs = await ai.retrieve({
      retriever: menuRetriever,
      query: input,
      options: { k: 3 },
    });

    // generate a response
   const { text } = await ai.generate({
      prompt: `
You are acting as a helpful AI assistant that can answer 
questions about the food available on the menu at Genkit Grub Pub.

Use only the context provided to answer the question.
If you don't know, do not make up an answer.
Do not add or change items on the menu.

Question: ${input}`,
      docs,
    });

    return text;
  }
);

独自のインデクサーとリトリーバーを作成する

リトリーバーは独自に作成することもできます。これは、ドキュメントが Genkit でサポートされていないドキュメント ストア（MySQL、Google ドライブなど）で管理されている場合に役立ちます。Genkit SDK には、ドキュメントの取得にカスタムコードを指定できる柔軟な方法が用意されています。Genkit の既存のリトリーバーの上にカスタム リトリーバーを定義し、高度な RAG 手法（再ランク付けやプロンプト拡張など）を追加で適用することもできます。

シンプルなリトリーバー

シンプルな Retriever を使用すると、既存のコードを Retriever に簡単に変換できます。

import { z } from "genkit";
import { searchEmails } from "./db";

ai.defineSimpleRetriever(
  {
    name: "myDatabase",
    configSchema: z
      .object({
        limit: z.number().optional(),
      })
      .optional(),
    // we'll extract "message" from the returned email item
    content: "message",
    // and several keys to use as metadata
    metadata: ["from", "to", "subject"],
  },
  async (query, config) => {
    const result = await searchEmails(query.text, { limit: config.limit });
    return result.data.emails;
  }
);

カスタム レトリーバ

import {
  CommonRetrieverOptionsSchema,
} from 'genkit/retriever';
import { z } from 'genkit';

export const menuRetriever = devLocalRetrieverRef('menuQA');

const advancedMenuRetrieverOptionsSchema = CommonRetrieverOptionsSchema.extend({
  preRerankK: z.number().max(1000),
});

const advancedMenuRetriever = ai.defineRetriever(
  {
    name: `custom/advancedMenuRetriever`,
    configSchema: advancedMenuRetrieverOptionsSchema,
  },
  async (input, options) => {
    const extendedPrompt = await extendPrompt(input);
    const docs = await ai.retrieve({
      retriever: menuRetriever,
      query: extendedPrompt,
      options: { k: options.preRerankK || 10 },
    });
    const rerankedDocs = await rerank(docs);
    return rerankedDocs.slice(0, options.k || 3);
  }
);

（extendPrompt と rerank は、フレームワークから提供されるものではなく、自分で実装する必要があります）。

その後、取得ツールを交換できます。

const docs = await ai.retrieve({
  retriever: advancedRetriever,
  query: input,
  options: { preRerankK: 7, k: 3 },
});

再ランクと 2 段階の取得

再ランキング モデル（クロスエンコーダとも呼ばれます）は、クエリとドキュメントが指定されたときに類似性スコアを出力するモデルの一種です。Google はこのスコアを使用して、クエリとの関連性に基づいてドキュメントの順序を並べ替えます。Reranker API は、ドキュメントのリスト（リトリーバーの出力など）を受け取り、クエリとの関連性に基づいてドキュメントを並べ替えます。このステップは、結果を微調整し、生成モデルに提供されるプロンプトで最も関連性の高い情報が使用されるようにするために役立ちます。

再ランク付けの例

Genkit の再ランクツールは、リトリーバーやインデクサーと同様の構文で定義されます。Genkit で再ランクツールを使用する例を次に示します。このフローでは、事前定義された Vertex AI 再ランクツールを使用して、指定されたクエリとの関連性に基づいて一連のドキュメントを再ランクします。

const FAKE_DOCUMENT_CONTENT = [
  'pythagorean theorem',
  'e=mc^2',
  'pi',
  'dinosaurs',
  'quantum mechanics',
  'pizza',
  'harry potter',
];

export const rerankFlow = ai.defineFlow(
  {
    name: 'rerankFlow',
    inputSchema: z.object({ query: z.string() }),
    outputSchema: z.array(
      z.object({
        text: z.string(),
        score: z.number(),
      })
    ),
  },
  async ({ query }) => {
    const documents = FAKE_DOCUMENT_CONTENT.map((text) =>
       ({ content: text })
    );

    const rerankedDocuments = await ai.rerank({
      reranker: 'vertexai/semantic-ranker-512',
      query:  ({ content: query }),
      documents,
    });

    return rerankedDocuments.map((doc) => ({
      text: doc.content,
      score: doc.metadata.score,
    }));
  }
);

この再ランク付けツールは、semantic-ranker-512 で Vertex AI genkit プラグインを使用してドキュメントのスコアとランク付けを行います。スコアが高いほど、ドキュメントとクエリの関連性が高くなります。

カスタム再ランクツール

特定のユースケースに合わせてカスタム再ランクを定義することもできます。これは、独自のカスタム ロジックまたはカスタムモデルを使用してドキュメントのランクを再調整する必要がある場合に便利です。カスタム再ランクを定義する簡単な例を次に示します。

export const customReranker = ai.defineReranker(
  {
    name: 'custom/reranker',
    configSchema: z.object({
      k: z.number().optional(),
    }),
  },
  async (query, documents, options) => {
    // Your custom reranking logic here
    const rerankedDocs = documents.map((doc) => {
      const score = Math.random(); // Assign random scores for demonstration
      return {
        ...doc,
        metadata: { ...doc.metadata, score },
      };
    });

    return rerankedDocs.sort((a, b) => b.metadata.score - a.metadata.score).slice(0, options.k || 3);
  }
);

定義したカスタム再ランクアルゴリズムは、RAG フロー内の他の再ランクアルゴリズムと同じように使用できるため、高度な再ランク戦略を柔軟に実装できます。