検索拡張生成(RAG)

Firebase Genkit には、検索拡張生成(RAG)フローの構築に役立つ抽象化と、関連ツールとの統合を提供するプラグインが用意されています。

RAG とは

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

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

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

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

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

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

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

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

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

インデクサ

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

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

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

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

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

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

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

エンベッダー

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

リトリーバー

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

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

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

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

  • Pinecone クラウド ベクトル データベース

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

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

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

RAG フローを定義する

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

依存関係のインストール

この例では、langchaingotextsplitter ライブラリと ledongthuc/pdf PDF 解析ライブラリを使用します。

go get github.com/tmc/langchaingo/textsplitter
go get github.com/ledongthuc/pdf

インデクサーを定義する

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

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

インデクサーを作成する

// Import Genkit's file-based vector retriever, (Don't use in production.)
import "github.com/firebase/genkit/go/plugins/localvec"

// Vertex AI provides the text-embedding-004 embedder model.
import "github.com/firebase/genkit/go/plugins/vertexai"
ctx := context.Background()

err := vertexai.Init(ctx, &vertexai.Config{})
if err != nil {
	log.Fatal(err)
}
err = localvec.Init()
if err != nil {
	log.Fatal(err)
}

menuPDFIndexer, _, err := localvec.DefineIndexerAndRetriever(
	"menuQA",
	localvec.Config{
		Embedder: vertexai.Embedder("text-embedding-004"),
	},
)
if err != nil {
	log.Fatal(err)
}

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

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

次の定義では、オーバーラップするチャンクが 20 文字で、200 文字のドキュメント セグメントを返すチャンク化関数を構成します。

splitter := textsplitter.NewRecursiveCharacter(
	textsplitter.WithChunkSize(200),
	textsplitter.WithChunkOverlap(20),
)

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

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

genkit.DefineFlow(
	"indexMenu",
	func(ctx context.Context, path string) (any, error) {
		// Extract plain text from the PDF. Wrap the logic in Run so it
		// appears as a step in your traces.
		pdfText, err := genkit.Run(ctx, "extract", func() (string, error) {
			return readPDF(path)
		})
		if err != nil {
			return nil, err
		}

		// Split the text into chunks. Wrap the logic in Run so it
		// appears as a step in your traces.
		docs, err := genkit.Run(ctx, "chunk", func() ([]*ai.Document, error) {
			chunks, err := splitter.SplitText(pdfText)
			if err != nil {
				return nil, err
			}

			var docs []*ai.Document
			for _, chunk := range chunks {
				docs = append(docs, ai.DocumentFromText(chunk, nil))
			}
			return docs, nil
		})
		if err != nil {
			return nil, err
		}

		// Add chunks to the index.
		err = ai.Index(ctx, menuPDFIndexer, ai.WithIndexerDocs(docs...))
		return nil, err
	},
)
// Helper function to extract plain text from a PDF. Excerpted from
// https://github.com/ledongthuc/pdf
func readPDF(path string) (string, error) {
	f, r, err := pdf.Open(path)
	if f != nil {
		defer f.Close()
	}
	if err != nil {
		return "", err
	}

	reader, err := r.GetPlainText()
	if err != nil {
		return "", err
	}

	bytes, err := io.ReadAll(reader)
	if err != nil {
		return "", err
	}
	return string(bytes), nil
}

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

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

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

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

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

	ctx := context.Background()

	err := vertexai.Init(ctx, &vertexai.Config{})
	if err != nil {
		log.Fatal(err)
	}
	err = localvec.Init()
	if err != nil {
		log.Fatal(err)
	}

	model := vertexai.Model("gemini-1.5-flash")

	_, menuPdfRetriever, err := localvec.DefineIndexerAndRetriever(
		"menuQA",
		localvec.Config{
			Embedder: vertexai.Embedder("text-embedding-004"),
		},
	)
	if err != nil {
		log.Fatal(err)
	}

	genkit.DefineFlow(
		"menuQA",
		func(ctx context.Context, question string) (string, error) {
			// Retrieve text relevant to the user's question.
			docs, err := menuPdfRetriever.Retrieve(ctx, &ai.RetrieverRequest{
				Document: ai.DocumentFromText(question, nil),
			})
			if err != nil {
				return "", err
			}

			// Construct a system message containing the menu excerpts you just
			// retrieved.
			menuInfo := ai.NewSystemTextMessage("Here's the menu context:")
			for _, doc := range docs.Documents {
				menuInfo.Content = append(menuInfo.Content, doc.Content...)
			}

			// Call Generate, including the menu information in your prompt.
			return ai.GenerateText(ctx, model,
				ai.WithMessages(
					ai.NewSystemTextMessage(`
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.`),
					menuInfo,
					ai.NewUserTextMessage(question)))
		})

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

リトリーバーは独自に作成することもできます。これは、ドキュメントが Genkit でサポートされていないドキュメント ストア(MySQL、Google ドライブなど)で管理されている場合に便利です。Genkit SDK には、ドキュメントの取得にカスタムコードを指定できる柔軟な方法が用意されています。

Genkit の既存のリトリーバーの上にカスタム リトリーバーを定義し、高度な RAG 手法(再ランク付けやプロンプト拡張など)を追加で適用することもできます。

たとえば、使用するカスタムの再ランク付け関数があるとします。次の例では、その関数を前述のメニュー リトリーバーに適用するカスタム リトリーバーを定義しています。

type CustomMenuRetrieverOptions struct {
	K          int
	PreRerankK int
}
advancedMenuRetriever := ai.DefineRetriever(
	"custom",
	"advancedMenuRetriever",
	func(ctx context.Context, req *ai.RetrieverRequest) (*ai.RetrieverResponse, error) {
		// Handle options passed using our custom type.
		opts, _ := req.Options.(CustomMenuRetrieverOptions)
		// Set fields to default values when either the field was undefined
		// or when req.Options is not a CustomMenuRetrieverOptions.
		if opts.K == 0 {
			opts.K = 3
		}
		if opts.PreRerankK == 0 {
			opts.PreRerankK = 10
		}

		// Call the retriever as in the simple case.
		response, err := menuPDFRetriever.Retrieve(ctx, &ai.RetrieverRequest{
			Document: req.Document,
			Options:  localvec.RetrieverOptions{K: opts.PreRerankK},
		})
		if err != nil {
			return nil, err
		}

		// Re-rank the returned documents using your custom function.
		rerankedDocs := rerank(response.Documents)
		response.Documents = rerankedDocs[:opts.K]

		return response, nil
	},
)