使用LangGraph迁移MapReduceDocumentsChain进行长文档的摘要

本文链接：https://blog.csdn.net/dgay_hua/article/details/146388736

在大数据处理和文本分析领域，MapReduce是一种非常重要的策略，用于处理和分析大型数据集。具体到文本处理方面，MapReduceDocumentsChain 구현了一种map-reduce策略，可以有效地处理长文本。本文将介绍如何从MapReduceDocumentsChain迁移到LangGraph，并探讨LangGraph在流处理、检查点恢复等方面的优势。

技术背景介绍

MapReduceDocumentsChain是一种特别适合长文本的处理策略，适用于需要将文本分割、并发处理然后整合结果的场景。通常的应用是文本摘要：将长文本分割为小段，然后对每段进行摘要，最后将所有摘要整合成一个最终的摘要。

核心原理解析

在MapReduce框架中，Map步骤通常是并行化的，即对每个文件段执行相同的处理步骤。Reduce步骤则负责整合Map步骤的结果。在文本摘要的上下文中，Map步骤生成各段文本的摘要，Reduce步骤则生成这些摘要的最终汇总。

LangGraph不仅支持MapReduce工作流，还提供了一些扩展功能，例如流式处理步骤，这样开发者可以更好地控制执行过程。LangGraph的检查点功能使得在处理失败时可以进行恢复，并且更容易融入对话式应用。

代码实现演示

下面我们将展示如何使用LangGraph实现MapReduceDocumentsChain的功能：

import operator
from typing import Annotated, List, TypedDict
from langchain_core.prompts import ChatPromptTemplate
from langgraph.constants import Send
from langgraph.graph import END, START, StateGraph

# 定义Map和Reduce的模板
map_template = "Write a concise summary of the following: {context}."
reduce_template = """
The following is a set of summaries:
{docs}
Take these and distill it into a final, consolidated summary
of the main themes.
"""

# 初始化模板
map_prompt = ChatPromptTemplate([("human", map_template)])
reduce_prompt = ChatPromptTemplate([("human", reduce_template)])

# 定义LangGraph的节点状态
class OverallState(TypedDict):
    contents: List[str]
    summaries: Annotated[list, operator.add]
    final_summary: str

class SummaryState(TypedDict):
    content: str

# 生成单个文档的摘要
async def generate_summary(state: SummaryState):
    response = await map_chain.ainvoke(state["content"])
    return {"summaries": [response]}

# 映射出文档摘要
def map_summaries(state: OverallState):
    return [
        Send("generate_summary", {"content": content}) for content in state["contents"]
    ]

# 生成最终摘要
async def generate_final_summary(state: OverallState):
    response = await reduce_chain.ainvoke(state["summaries"])
    return {"final_summary": response}

# 构建并编译图
graph = StateGraph(OverallState)
graph.add_node("generate_summary", generate_summary)
graph.add_node("generate_final_summary", generate_final_summary)
graph.add_conditional_edges(START, map_summaries, ["generate_summary"])
graph.add_edge("generate_summary", "generate_final_summary")
graph.add_edge("generate_final_summary", END)
app = graph.compile()

# 以流模式调用图
async for step in app.astream({"contents": [doc.page_content for doc in documents]}):
    print(step)