检索增强生成 (RAG)

Firebase Genkit 提供了可帮助您构建检索增强生成 (RAG) flow 的抽象,以及可提供与相关工具集成的插件。

什么是 RAG?

检索增强生成是一种技术,用于将外部信息来源整合到 LLM 的回答中。能够做到这一点非常重要,因为虽然 LLM 通常是根据大量材料进行训练的,但实际使用 LLM 通常需要特定的领域知识(例如,您可能希望使用 LLM 回答客户有关贵公司产品的问题)。

一种解决方案是使用更具体的数据对模型进行微调。但是,就计算费用和准备充足训练数据所需的工作量而言,这可能都很昂贵。

相比之下,RAG 的工作原理是,在将外部数据源传递给模型时将其整合到提示中。例如,您可以想象一下,通过在前面添加一些相关信息,可以对提示“What is Bart's relationship to Lisa?”这一提示进行扩展(“增强”),从而生成提示“Homer and Marge's children are named Bart, Lisa, and Maggie.What is Bart's relationship to Lisa?”

这样做具有很多优势:

  • 这可能更具成本效益,因为您无需重新训练模型。
  • 您可以持续更新数据源,并且 LLM 可以立即利用更新后的信息。
  • 现在,您可以在 LLM 的回答中引用参考信息。

另一方面,使用 RAG 本质上意味着较长的提示,并且某些 LLM API 服务会针对您发送的每个输入词元收费。归根结底,您必须评估应用的费用权衡。

RAG 是一个非常广泛的领域,有许多不同的技术可用于实现最佳质量的 RAG。核心 Genkit 框架提供了两个主要抽象,以帮助您执行 RAG:

  • 索引器:将文档添加到“索引”中。
  • 嵌入器:将文档转换为矢量表示
  • 检索器:根据查询从“索引”中检索文档。

这些定义本来就很宽泛,因为 Genkit 对“索引”是什么或者如何从其中检索文档没有确切的看法。Genkit 仅提供 Document 格式,而其他所有内容均由检索器或索引器实现提供程序定义。

索引器

索引负责跟踪您的文档,以便您在进行特定查询时快速检索相关文档。这通常是通过矢量数据库来实现的,矢量数据库使用称为嵌入的多维矢量将文档编入索引。文本嵌入(不透明)表示由一段文本表示的概念;这些概念是使用特殊用途的机器学习模型生成的。通过使用嵌入将文本编入索引,矢量数据库能够对概念相关的文本进行聚类,并检索与新型文本字符串(查询)相关的文档。

您需要先将文档注入到文档索引中,然后才能检索文档以进行生成。典型的注入 flow 会执行以下操作:

  1. 将大型文档拆分成较小的文档,以便仅使用相关部分来增强提示,即“分块”。这一点很有必要,因为许多 LLM 的上下文窗口有限,因此在提示中包含整个文档是不切实际的。

    Genkit 不提供内置分块库;不过,有可用的开源库与 Genkit 兼容。

  2. 为每个分块生成嵌入。根据您使用的数据库,您可以使用嵌入生成模型明确执行此操作,也可以使用数据库提供的嵌入生成器。

  3. 将文本块及其索引添加到数据库中。

如果您处理的是稳定的数据源,则可以不经常或仅运行一次注入 flow。另一方面,如果您处理的是经常更改的数据,则可以持续运行注入 flow(例如,在 Cloud Firestore 触发器中,每当文档更新时)。

嵌入器

嵌入器是一种函数,它会接受内容(文本、图片、音频等)并创建一个数字矢量,以对原始内容的语义含义进行编码。如上所述,嵌入器会作为索引编入过程的一部分来使用,但也可以独立使用来创建没有索引的嵌入。

检索器

检索器是一个概念,用于封装与任何类型的文档检索相关的逻辑。最常见的检索情况通常包括从矢量存储区检索,但在 Genkit 中,检索器可以是任何返回数据的函数。

如需创建检索器,您可以使用提供的某个实现,也可以创建自己的实现。

支持的索引器、检索器和嵌入器

Genkit 通过其插件系统提供索引器和检索器支持。以下插件受官方支持:

此外,Genkit 还通过预定义的代码模板支持以下矢量存储区,您可以根据自己的数据库配置和架构对其进行自定义:

  • 将 PostgreSQL 与 pgvector 搭配使用

嵌入模型支持通过以下插件提供:

插件 模型
Google 生成式 AI Gecko 文本嵌入
Google Vertex AI Gecko 文本嵌入

定义 RAG Flow

以下示例展示了如何将餐厅菜单 PDF 文档集合注入到矢量数据库中,并检索这些文档,以便在确定可提供食材的 flow 中使用。

安装依赖项

在此示例中,我们将使用 langchaingo 中的 textsplitter 库和 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 库,该库提供了一个简单的文本拆分器,用于将文档拆分为可矢量化的片段。

以下定义会将分块函数配置为返回 200 个字符的文档段,并且分块之间有 20 个字符的重叠。

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

如需了解此库的更多分块选项,请参阅 langchaingo 文档

定义索引器 flow

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
}

运行索引器 flow

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

运行 indexMenu flow 后,矢量数据库将添加文档种子,并准备好在包含检索步骤的 Genkit flow 中使用。

定义包含检索的 flow

以下示例展示了如何在 RAG flow 中使用检索器。与索引器示例一样,此示例使用 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
	},
)