Graph RAG 生产部署指南 — 一步一步详解

图片由 JJ Ying 在 Unsplash 提供

在概念性地讨论了Graph RAG之后，让我们将其投入生产。这是如何将GraphRAG投入生产的方法：完全采用无服务器架构，完全并行化以最小化推理和索引时间，并且永远不需要接触图数据库（保证！）。

在本文中，我将向您介绍graphrag-lite，这是一个端到端的Graph RAG数据摄入和查询实现。我将graphrag-lite作为一个开源项目发布，以使您在GCP上部署graphrag时更加轻松。graphrag-lite是专门为Google Cloud设计的，开箱即用。代码采用模块化设计，可根据您的选择进行调整。

检索增强生成本身尚未描述任何特定的架构或方法。它仅描述了使用任意检索方法增强给定生成任务的过程。原始的RAG论文（由Lewis等人撰写的《检索增强生成在知识密集型NLP任务中的应用》（Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks））比较了双塔嵌入方法与词袋检索方法。

现代的问答系统能够区分本地和全局问题。在一个非结构化的知识库中，一个本地（提取型）问题可能是“2023年诺贝尔和平奖得主是谁？”而一个全局（聚合型）问题可能是“你知道最近的诺贝尔奖得主有哪些？”text2embedding RAG在处理全局和结构化问题时存在明显不足。而Graph RAG可以弥补这些不足，并且做得很好！通过一个抽象层，它学习了知识图谱社区的语义。这构建了一个对索引数据集更“全局”的理解。这里有一篇概念性介绍Graph RAG的文章可以阅读。

一个 Graph RAG 管道通常会遵循以下步骤：

这是主要的导入步骤。您的LLM会使用一个提示来扫描每一份传入的文档，以提取我们知识图谱的相关节点和边。您需要多次迭代这个提示，以确保捕获所有相关信息。

你将提取的节点和边存储在你选择的数据存储中。专用的图数据库是一个选项，但它们通常比较繁琐。Graph2nosql 是一个基于 Python 的接口，用于在 Firestore 或任何其他 NoSQL 数据库中存储和管理知识图谱。我开源了这个项目，因为我没有在市场上找到任何其他类似的知识图谱原生选项。

一旦你存储了知识图谱数据，你将使用社区检测算法来识别那些在彼此之间比在图的其余部分更紧密连接的节点组。在知识图谱的背景下，假设密集的社区涵盖了共同的主题。

然后你指示你的LLM为每个图社区生成一份报告。这些社区报告有助于抽象出单一主题，从而把握数据集中的更广泛、全局的概念。社区报告将与你的知识图谱一起存储。这标志着管道的摄入层的结束。

在查询时，你遵循 map-reduce 模式为知识图谱中的每个社区报告生成用户查询的中间响应。你让大语言模型（LLM）评估每个中间查询响应的相关性。最后，你根据相关性对中间响应进行排序，并选择前 n 个作为最终响应给用户的上下文。

Graph RAG 分步逻辑 — 作者作品图

在初始的导入步骤中，你指示你的LLM将输入文档编码为一个图。一个详细的提示要求你的LLM首先识别给定类型的节点，然后识别这些节点之间的边。就像任何LLM提示一样，这个问题没有一个固定的解决方案。这里是我基于Microsoft的开源实现构建的图提取提示的核心内容：

    -目标-  
    给定一个可能与此活动相关的文本文档和一组实体类型，识别文本中的所有这些类型的实体以及这些实体之间的所有关系。  

    -步骤-  

1. 识别所有实体。对于每个识别出的实体，提取以下信息：  
    - entity_name: 实体名称，首字母大写  
    - entity_type: 以下类型之一：[{entity_types}]  
    - entity_description: 实体属性和活动的全面描述  
    格式化每个实体为 ("entity"{tuple_delimiter}<entity_name>{tuple_delimiter}<entity_type>{tuple_delimiter}<entity_description>  

2. 从步骤1中识别出的实体中，识别所有*明显相关*的实体对（source_entity, target_entity）。  
    对于每对相关实体，提取以下信息：  
    - source_entity: 步骤1中识别出的源实体名称  
    - target_entity: 步骤1中识别出的目标实体名称  
    - relationship_description: 说明为什么你认为源实体和目标实体彼此相关  
    - relationship_strength: 表示源实体和目标实体之间关系强度的数值评分  
     格式化每个关系为 ("relationship"{tuple_delimiter}<source_entity>{tuple_delimiter}<target_entity>{tuple_delimiter}<relationship_description>{tuple_delimiter}<relationship_strength>)  

3. 将步骤1和2中识别出的所有实体和关系作为单个列表返回，使用**{record_delimiter}**作为列表分隔符。  

4. 完成后，输出 {completion_delimiter}  

    <多轮示例>  

    -真实数据-  
    ######################  
    Entity_types: {entity_types}  
    Text: {input_text}  
    ######################  
    输出：

提取步骤负责确定哪些信息将反映在您的知识库中。因此，您应该使用一个较为强大的模型，例如 Gemini 1.5 Pro。您可以通过使用多轮版本的 Gemini 1.5 Pro 来进一步提高结果的鲁棒性，并多次查询模型以改进其结果。以下是我如何在 graphrag-lite 中实现图提取循环的：

    class GraphExtractor:  
        def __init__(self, graph_db) -> None:  
            self.tuple_delimiter = "<|>"  
            self.record_delimiter = "##"  
            self.completion_delimiter = ""  
            self.entity_types = ["organization", "person", "geo", "event"]  

            self.graph_extraction_system = prompts.GRAPH_EXTRACTION_SYSTEM.format(  
                entity_types=", ".join(self.entity_types),  
                record_delimiter=self.record_delimiter,  
                tuple_delimiter=self.tuple_delimiter,  
                completion_delimiter=self.completion_delimiter,  
            )  

            self.llm = LLMSession(system_message=self.graph_extraction_system,  
                                  model_name="gemini-1.5-pro-001")  

        def __call__(self, text_input: str, max_extr_rounds: int = 5) -> None:  

            input_prompt = self._construct_extractor_input(input_text=text_input)  

            print("+++++ 初始化图提取 +++++")  

            init_extr_result = self.llm.generate_chat(  
                client_query_string=input_prompt, temperature=0, top_p=0)  
            print(f"初始结果: {init_extr_result}")  

            for round_i in range(max_extr_rounds):  

                print(f"+++++ 继续图提取第 {round_i} 轮 +++++")  

                round_response = self.llm.generate_chat(  
                    client_query_string=prompts.CONTINUE_PROMPT, temperature=0, top_p=0)  
                init_extr_result += round_response or ""  

                print(f"本轮响应: {round_response}")  

                if round_i >= max_extr_rounds - 1:  
                    break  

                completion_check = self.llm.generate_chat(  
                    client_query_string=prompts.LOOP_PROMPT, temperature=0, top_p=0)  

                if "YES" not in completion_check:  
                    print(  
                        f"+++++ 在第 {round_i} 轮后完成，完成检查 +++++")  
                    break

首先，我调用多轮模型来提取节点和边。然后，我多次请求模型来改进之前的提取结果。

在 graphrag-lite 实现中，提取模型的调用是由前端客户端发起的。如果你想减轻客户端的负载，可以将提取查询外包给一个微服务。

一旦从文档中提取出节点和边，你需要将它们存储在一个可访问的格式中。图数据库是一种选择，但它们也可能比较繁琐。对于你的知识图谱，你可能在寻找一种更轻量级的解决方案。我也有同感，因为我没有找到任何开源的知识图谱原生库，所以我开源了 graph2nosql。Graph2nosql 是一个简单的知识图谱原生 Python 接口。它帮助你在任何 NoSQL 数据库中存储和管理你的知识图谱，而无需在技术栈中添加图数据库或学习 Cypher。

Graph2nosql 是为知识图谱检索而设计的，考虑到了图谱检索（Graph RAG）。该库围绕三种主要数据类型设计：EdgeData、NodeData 和 CommunityData。节点通过 uid 标识。边通过源节点和目标节点的 uid 以及边的 uid 标识。由于 uids 可以自由设计，Graph2nosql 数据模型为任何规模的知识图谱留出了空间。您甚至可以添加文本或图嵌入。这允许基于嵌入的分析、边预测以及额外的文本嵌入检索（考虑混合 RAG）。

Graph2nosql 原生设计围绕着 Firestore。

    @dataclass  
    class EdgeData:  
        源节点ID: str  
        目标节点ID: str   
        描述: str  
        边ID: str | None = None  
        文档ID: str | None = None  

    @dataclass  
    class NodeData:  
        节点ID: str  
        节点标题: str  
        节点类型: str  
        节点描述: str  
        节点度数: int  
        文档ID: str   
        社区ID: int | None = None # 基于源文档的社区ID   
        出边: list[str] = field(default_factory=list)  
        入边: list[str] = field(default_factory=list)  # 在有向图的情况下  
        嵌入: list[float] = field(default_factory=list)  # 表示节点的文本嵌入，例如标题和描述的组合  

    @dataclass  
    class CommunityData:  
        标题: str # 社区标题，如果尚未计算则为 None  
        社区节点: set[str] = field(default_factory=set) # 属于社区的节点ID列表  
        摘要: str | None = None # 社区描述，如果尚未计算则为 None  
        文档ID: str | None = None # 该实体的源知识库文档标识符  
        社区ID: str | None = None # 社区标识符  
        社区嵌入: Tuple[float, ...] = field(default_factory=tuple) # 表示社区的文本嵌入  
        评分: int | None = None  
        评分解释: str | None = None  
        发现: list[dict] | None = None

要通过 graph2nosql 存储你的图数据，只需在解析提取步骤的结果时运行以下代码。这里是 graphrag-lite 的实现。

from graph2nosql.graph2nosql.graph2nosql import NoSQLKnowledgeGraph  
from graph2nosql.databases.firestore_kg import FirestoreKG  
from graph2nosql.datamodel import data_model  

fskg = FirestoreKG(  
        gcp_project_id=project_id,  
        gcp_credential_file=firestore_credential_file,  
        firestore_db_id=database_id,  
        node_collection_id=node_coll_id,  
        edges_collection_id=edges_coll_id,  
        community_collection_id=community_coll_id)  

node_data = data_model.NodeData(  
        node_uid=entity_name,  
        node_title=entity_name,  
        node_type=entity_type,  
        node_description=entity_description,  
        document_id=str(source_doc_id),  
        node_degree=0)  

fskg.add_node(node_uid=entity_name, node_data=node_data)

在将所有相关节点和边存储到您的图数据库后，您可以开始构建抽象层。一种方法是找到描述相似概念的节点，并描述它们的语义连接方式。Graph2nosql 提供了内置的社区检测功能，例如基于 Louvain 社区的方法。

根据你的提取结果质量，你可能会在知识图谱中发现零度节点。从经验来看，零度节点往往是重复项。GraphRAG-lite 使用图社区作为主要抽象层，因此你应该删除没有边的节点。因此，考虑另一个重复项/合并步骤和/或基于描述和图嵌入的节点预测步骤，以添加提取步骤中可能遗漏的边是有意义的。在 GraphRAG-lite 中，我目前只是简单地删除所有零度节点。

    # 清除所有没有边的节点  
    fskg.clean_zerodegree_nodes()  

    # 基于清理后的图生成社区  
    comms = kg.get_louvain_communities()

这里提供了社区检测的 graphrag-lite 实现。

上述提到的GraphRAG管道在每份文档摄入和用户查询时会进行多次LLM调用。例如，为了为每个新索引的文档生成多个社区报告，或者在查询时为多个社区生成中间响应。如果同时处理，会导致糟糕的用户体验。特别是在大规模情况下，用户可能需要等待几分钟到几小时才能收到查询的响应。幸运的是，如果你以正确的方式设计LLM提示，可以将它们设计为“无状态工作者”。无状态处理架构的优势是双重的。首先，它们易于并行化。其次，它们易于作为无服务器基础设施实现。结合并行化和无服务器架构，可以最大化吞吐量的可扩展性，并最小化空闲集群设置的成本。

在 graphrag-lite 架构中，我将社区报告生成和中间查询生成作为无服务器 Cloud Run 微服务工作者进行托管。这些工作者通过 GCP 的无服务器消息队列 PubSub 接收消息。

graphrag-lite 的无服务器和分布式摄入及查询管道 — 作者供图

运行社区检测后，你现在知道了多个社区成员节点集合。每个集合代表了知识图谱中的一个语义主题。社区报告步骤需要对这些源自知识库中不同文档的概念进行抽象。我再次基于微软的实现，并添加了一个易于解析的结构化输出的函数调用。

    你是一个AI助手，帮助人类分析师进行一般信息发现。信息发现是指识别和评估与特定实体（例如组织和个人）相关的相关信息的过程。  

    # 目标  
    根据一个社区的实体列表及其关系和可选的相关声明，撰写一份全面的报告。该报告将用于向决策者提供与社区及其潜在影响相关的信息。报告内容包括社区的关键实体概述、其法律合规性、技术能力、声誉以及值得注意的声明。  

    # 报告结构  

    报告应包括以下部分：  

    - 标题：代表社区关键实体的社区名称 - 标题应简短但具体。如果可能，标题中应包含代表性的命名实体。  
    - 摘要：社区整体结构的执行摘要，其实体之间的关系，以及与其实体相关的重大信息。  
    - 影响严重性评分：介于0-10之间的浮点分数，表示社区内实体造成的影响严重性。影响是指社区的重要性评分。  
    - 评分解释：给出影响严重性评分的单句解释。  
    - 详细发现：关于社区的5-10个关键见解。每个见解应包括简短的摘要，然后是多个段落的解释性文本，根据下面的接地规则进行说明。要全面。

社区报告生成也展示了知识图谱检索方面最大的挑战。理论上，任何文档都可以向图中的每个现有社区添加一个新的节点。在最坏的情况下，每当添加一个新的文档时，你都需要重新生成知识库中的每个社区报告。实际上，重要的是要包含一个检测步骤，以识别文档上传后哪些社区发生了变化，从而仅针对调整后的社区生成新的报告。

由于每次上传文档时都需要重新生成多个社区报告，我们在并发运行这些请求时也面临着显著的延迟挑战。因此，你应该将这些工作外包并并行化，使用异步工作者来处理。如前所述，graphrag-lite 使用无服务器架构解决了这个问题。我使用 PubSub 作为消息队列来管理任务并确保处理。Cloud Run 作为计算平台托管无状态工作者，调用 LLM 进行处理。生成时，他们使用如上所示的提示。

这里是用于社区报告生成的状态无畏工作者中运行的代码:

    def async_generate_comm_report(self, comm_members: set[str]) -> data_model.CommunityData:  

            llm = LLMSession(system_message=prompts.COMMUNITY_REPORT_SYSTEM,  
                             model_name="gemini-1.5-flash-001")  

            response_schema = {  
                "type": "object",  
                "properties": {  
                        "title": {  
                            "type": "string"  
                        },  
                    "summary": {  
                            "type": "string"  
                            },  
                    "rating": {  
                            "type": "int"  
                            },  
                    "rating_explanation": {  
                            "type": "string"  
                            },  
                    "findings": {  
                            "type": "array",  
                            "items": {  
                                "type": "object",  
                                "properties": {  
                                    "summary": {  
                                        "type": "string"  
                                    },  
                                    "explanation": {  
                                        "type": "string"  
                                    }  
                                },  
                                # 确保每个发现都包含这两个字段  
                                "required": ["summary", "explanation"]  
                            }  
                            }  
                },  
                # 列出顶级所需的字段  
                "required": ["title", "summary", "rating", "rating_explanation", "findings"]  
            }  

            comm_report  = llm.generate(client_query_string=prompts.COMMUNITY_REPORT_QUERY.format(  
                entities=comm_nodes,  
                relationships=comm_edges,  
                response_mime_type="application/json",  
                response_schema=response_schema  
            ))  

    comm_data = data_model.CommunityData(title=comm_report_dict["title"],                                              summary=comm_report_dict["summary"],                                                rating=comm_report_dict["rating"],             rating_explanation=comm_report_dict["rating_explanation"],               findings=comm_report_dict["findings"],  
    community_nodes=comm_members)  

     return comm_data

这完成了数据摄入管道。

最后，你到了查询时间。为了生成最终的用户响应，你需要生成一系列中间响应（每个社区报告一个）。每个中间响应将用户查询和一个社区报告作为输入。然后，你需要根据相关性对这些中间响应进行评分。最后，你使用最相关的社区报告以及相关成员节点的节点描述等额外信息作为最终查询上下文。在大规模社区报告数量较多的情况下，这再次带来了延迟和成本的挑战。与之前一样，你也应该将中间响应生成（映射步骤）并行化处理，以跨无服务器微服务进行。未来，你可以通过添加一个过滤层来预判社区报告对用户查询的相关性，从而显著提高效率。

映射步骤的微服务如下所示:

    def 生成响应(client_query: str, community_report: dict):  

        llm = LLMSession(  
            system_message=MAP_SYSTEM_PROMPT,  
            model_name="gemini-1.5-pro-001"  
        )  

        响应模式 = {  
            "type": "object",  
            "properties": {  
                "响应": {  
                    "type": "string",  
                    "description": "用户问题的原始字符串响应。",  
                },  
                "得分": {  
                    "type": "number",  
                    "description": "给定社区报告上下文对回答用户问题的相关性得分 [0.0, 10.0]",  
                },  
            },  
            "required": ["响应", "得分"],  
        }  

        查询提示 = MAP_QUERY_PROMPT.format(  
            context_community_report=community_report, user_question=client_query)  

        响应 = llm.generate(client_query_string=查询提示,  
                     response_schema=响应模式,  
                     response_mime_type="application/json")  

        return 响应

地图步骤微服务使用了以下提示：

    ---角色---  
    你是一个专家代理，根据组织成知识图谱的上下文回答问题。  
    你将被提供一个从同一知识图谱中提取的社区报告。  

    ---目标---  
    生成一个包含关键点列表的响应，以回答用户的问题，总结给定社区报告中的所有相关信息。  

    你应该仅使用下面的社区描述中的数据作为生成响应的上下文。  
    如果你不知道答案，或者输入的社区描述中没有足够的信息来提供答案，请回答“根据给定的社区上下文，无法回答用户的问题。”。  

    你的响应应始终包含以下元素：  
    - 查询响应：基于提供的上下文，对给定用户查询的全面且真实的响应。  
    - 重要性评分：一个介于0-10之间的整数评分，表示该点在回答用户问题中的重要性。对于“我不知道”的类型响应，评分应为0。  

    响应应格式化为以下JSON格式：  
    {{"response": "描述点1 [数据：报告（报告ID）]", "score": 评分值}}  

    ---上下文社区报告---  
    {context_community_report}  

    ---用户问题---  
    {user_question}  

    ---JSON响应---  
    格式化为以下JSON响应：  
    {{"response": "描述点1 [数据：报告（报告ID）]", "score": 评分值}}  

    响应： 

为了成功执行reduce步骤，你需要存储中间响应，以便在查询时访问。在graphrag-lite中，我使用Firestore作为微服务之间的共享状态。触发中间响应生成后，客户端还会定期检查共享状态中是否存在所有预期的条目。以下是从graphrag-lite中提取的代码片段，展示了我是如何将每个社区报告提交到PubSub队列的。之后，我会定期查询共享状态，检查所有中间响应是否已处理完毕。最后，使用得分最高的社区报告作为上下文生成最终的用户响应。

    class KGraphGlobalQuery:  
        def __init__(self) -> None:  
            # 使用消息队列、知识图谱、共享NoSQL状态初始化  
            pass  

        @observe()  
        def __call__(self, user_query: str) -> str:  

            # 方法将自然语言用户查询转换为最终答案并返回给客户端  
            comm_report_list = self._get_comm_reports()  

            # 将用户查询与现有社区报告配对  
            query_msg_list = self._context_builder(  
                user_query=user_query, comm_report_list=comm_report_list)  

            # 将配对发送到PubSub队列进行任务调度  
            for msg in query_msg_list:  
                self._send_to_mq(message=msg)  
            print("中间响应请求已发送到消息队列")  

            # 定期查询共享状态以检查处理完成情况并获取中间响应  
            intermediate_response_list = self._check_shared_state(  
                user_query=user_query)  

            # 根据有用性构建最终上下文  
            sorted_final_responses = self._filter_and_sort_responses(intermediate_response_list=intermediate_response_list)  

            # 获取选定社区的完整社区报告  
            comm_report_list = self._get_communities_reports(sorted_final_responses)  

            # 根据最终上下文和社区报告生成并返回最终响应  
            final_response_system = prompts.GLOBAL_SEARCH_REDUCE_SYSTEM.format(  
                response_type="详细且全面的学术风格分析，至少包含8-10句话，跨越2-3段落。")  

            llm = LLMSession(  
                system_message=final_response_system,  
                model_name="gemini-1.5-pro-001"  
            )  

            final_query_string = prompts.GLOBAL_SEARCH_REDUCE_QUERY.format(  
                report_data=comm_report_list,  
                user_query=user_query  
            )  
            final_response = llm.generate(client_query_string=final_query_string)  
            return final_response

一旦找到所有条目，客户端将根据选定的社区上下文触发最终的用户响应生成。

Graph RAG 是每个 ML 工程师都应该添加到他们工具箱中的强大技术。每一个问答类型的应用最终都会到达这样一个点：纯粹的提取式、“局部”查询已经不够用了。现在，通过 graphrag-lite，你有了一个轻量级、云原生且无服务器的实现，可以快速复制。

尽管有这些优势，需要注意的是，目前状态下Graph RAG仍然比text2emb RAG消耗更多的LLM输入令牌。这通常会导致查询和文档索引的延迟和成本显著增加。然而，经历了结果质量的提升后，我确信在正确的应用场景中，Graph RAG是值得投入时间和金钱的。

RAG 应用程序最终将走向混合方向。提取式查询可以通过 text2emb RAG 高效且准确地处理。全局抽象查询可能需要知识图谱作为替代的检索层。最后，这两种方法在处理定量和分析查询时表现不佳。因此，增加一个第三层的 text2sql 检索层将带来巨大的价值。为了完善这一图景，用户查询可以首先在三种检索方法之间进行分类。这样，每个查询都可以最高效地获得适当的信息量和深度。

我迫不及待地想看看这还会走向何方。你一直在研究哪些其他的检索方法？