构建您自己的 RAG 助手

FlexRAG 提供了一套灵活的模块化组件用于构建 RAG 助手。您可以通过继承 Assistant 类并将其注册到 ASSISTANTS 装饰器中来构建您自己的 RAG 助手。

定义您的 RAG 助手

为了构建您的 RAG 助手，您可以创建一个 Python 脚本文件并导入必要的 FlexRAG 模块。下面是一个如何构建 RAG 助手的示例。在示例中，我们通过继承 AssistantBase 类定义了一个名为 SimpleAssistant 的 RAG 助手，该助手包含了一个检索器 (FlexRetriever) 和一个生成器 (OpenAIGenerator)。每当用户提出一个问题时，SimpleAssistant 都使用 FlexRetriever 在数据库中检索相关文档，随后将相关文档拼接到 prompt 中，使用 OpenAIGenerator 来生成最终回复。

from dataclasses import dataclass

from flexrag.assistant import ASSISTANTS, AssistantBase
from flexrag.models import OpenAIGenerator, OpenAIGeneratorConfig
from flexrag.prompt import ChatPrompt, ChatTurn
from flexrag.retriever import FlexRetriever, FlexRetrieverConfig


@dataclass
class SimpleAssistantConfig(FlexRetrieverConfig, OpenAIGeneratorConfig): ...


@ASSISTANTS("simple", config_class=SimpleAssistantConfig)
class SimpleAssistant(AssistantBase):
    def __init__(self, config: SimpleAssistantConfig):
        self.retriever = FlexRetriever(config)
        self.generator = OpenAIGenerator(config)
        return

    def answer(self, question: str) -> str:
        prompt = ChatPrompt()
        context = self.retriever.search(question)[0]
        prompt_str = "Please answer the following question based on the given text.\n\n"
        prompt_str += f"Question: {question}\n\n"
        for n, ctx in enumerate(context):
            prompt_str += f"Context {n}: {ctx.data['text']}\n"
        prompt.update(ChatTurn(role="user", content=prompt_str))
        response = self.generator.chat([prompt])[0][0]
        prompt.update(ChatTurn(role="assistant", content=response))
        return response

评估您的 RAG 应用

在完成 SimpleAssistant 的定义后，您仅需在命令中添加 user_module=<your_module_path>就可以使用 FlexRAG 内置的工具来评估您的 RAG 助手。

举例来说，您可以通过下面的命令在 Natural Questions 数据集上评估您的 RAG 助手：

DB_PATH=<path_to_database>
OPENAI_KEY=<your_openai_key>
MODULE_PATH=<path_to_simple_assistant_module>

python -m flexrag.entrypoints.eval_assistant \
    user_module=${MODULE_PATH} \
    name=nq \
    split=test \
    assistant_type=simple \
    simple_config.model_name='gpt-4o-mini' \
    simple_config.api_key=${OPENAI_KEY} \
    simple_config.retriever_path=${DB_PATH} \
    simple_config.used_indexes=[contriever] \
    eval_config.metrics_type=[retrieval_success_rate,generation_f1,generation_em] \
    eval_config.retrieval_success_rate_config.eval_field=text \
    log_interval=10

在仓库 FlexRAG_Examples 中，我们提供了一些详细的示例来展示如何使用 FlexRAG 来构建您的 RAG 助手。您可以访问该仓库以学习更多有关如何构建 RAG 助手的知识。