237 lines
7.4 KiB
Python
237 lines
7.4 KiB
Python
"""LangGraph 并行节点与 Reducer 示例。
|
||
|
||
本例包含两个图：
|
||
|
||
1. good_graph：两个搜索节点并行执行，并用 reducer 合并结果。
|
||
2. bad_graph：两个并行节点同时覆盖同一字段，触发 InvalidUpdateError。
|
||
|
||
主图结构：
|
||
|
||
┌─> search_web ─────┐
|
||
START -> prepare ──┤ ├─> summarize -> END
|
||
└─> search_database ┘
|
||
|
||
运行：
|
||
python 10_parallel_and_reducers.py
|
||
"""
|
||
|
||
import operator
|
||
import time
|
||
from typing import Annotated, TypedDict
|
||
|
||
from langgraph.errors import InvalidUpdateError
|
||
from langgraph.graph import END, START, StateGraph
|
||
|
||
|
||
# -----------------------------------------------------------------------------
|
||
# 1. Reducer
|
||
# -----------------------------------------------------------------------------
|
||
|
||
|
||
def merge_unique(left: list[str], right: list[str]) -> list[str]:
|
||
"""合并两个列表并去重，同时保留第一次出现的顺序。"""
|
||
return list(dict.fromkeys(left + right))
|
||
|
||
|
||
class ResearchState(TypedDict):
|
||
topic: str
|
||
|
||
# operator.add：把各节点返回的列表连接起来。
|
||
results: Annotated[list[str], operator.add]
|
||
|
||
# 自定义 reducer：合并日志并去除重复内容。
|
||
logs: Annotated[list[str], merge_unique]
|
||
|
||
# 没有 reducer：使用默认覆盖策略。
|
||
# 这里只有 summarize 节点更新它，因此不会发生并行冲突。
|
||
summary: str
|
||
|
||
|
||
# -----------------------------------------------------------------------------
|
||
# 2. 主图节点
|
||
# -----------------------------------------------------------------------------
|
||
|
||
|
||
def prepare(state: ResearchState) -> dict:
|
||
"""准备搜索；完成后两个下游节点会在同一 super-step 并行执行。"""
|
||
print(f"[prepare] 准备研究：{state['topic']}")
|
||
return {"logs": ["准备完成"]}
|
||
|
||
|
||
def search_web(state: ResearchState) -> dict:
|
||
"""模拟耗时 1 秒的网页搜索。"""
|
||
print("[search_web] 开始")
|
||
time.sleep(1)
|
||
print("[search_web] 完成")
|
||
return {
|
||
"results": [
|
||
f"网页结果：{state['topic']} 支持有状态工作流",
|
||
"共同结果：LangGraph 使用 StateGraph 编排节点",
|
||
],
|
||
"logs": ["网页搜索完成", "搜索阶段完成"],
|
||
}
|
||
|
||
|
||
def search_database(state: ResearchState) -> dict:
|
||
"""模拟耗时 1 秒的数据库搜索。"""
|
||
print("[search_database] 开始")
|
||
time.sleep(1)
|
||
print("[search_database] 完成")
|
||
return {
|
||
"results": [
|
||
f"数据库结果：{state['topic']} 支持持久化和恢复",
|
||
"共同结果：LangGraph 使用 StateGraph 编排节点",
|
||
],
|
||
"logs": ["数据库搜索完成", "搜索阶段完成"],
|
||
}
|
||
|
||
|
||
def summarize(state: ResearchState) -> dict:
|
||
"""汇合节点：两个搜索分支都结束后才会执行。"""
|
||
print(f"[summarize] 收到 {len(state['results'])} 条结果")
|
||
|
||
lines = "\n".join(
|
||
f"{index}. {result}"
|
||
for index, result in enumerate(state["results"], start=1)
|
||
)
|
||
return {
|
||
"summary": f"主题：{state['topic']}\n{lines}",
|
||
"logs": ["摘要生成完成"],
|
||
}
|
||
|
||
|
||
# -----------------------------------------------------------------------------
|
||
# 3. 构建正确的并行图
|
||
# -----------------------------------------------------------------------------
|
||
|
||
|
||
good_builder = StateGraph(ResearchState)
|
||
good_builder.add_node("prepare", prepare)
|
||
good_builder.add_node("search_web", search_web)
|
||
good_builder.add_node("search_database", search_database)
|
||
good_builder.add_node("summarize", summarize)
|
||
|
||
good_builder.add_edge(START, "prepare")
|
||
|
||
# 一个节点拥有两条普通出边：下一 super-step 中两个目标节点并行执行。
|
||
good_builder.add_edge("prepare", "search_web")
|
||
good_builder.add_edge("prepare", "search_database")
|
||
|
||
# 两个分支都指向 summarize。LangGraph 会等待同一轮的两个分支完成后汇合。
|
||
good_builder.add_edge("search_web", "summarize")
|
||
good_builder.add_edge("search_database", "summarize")
|
||
good_builder.add_edge("summarize", END)
|
||
|
||
good_graph = good_builder.compile()
|
||
|
||
|
||
# -----------------------------------------------------------------------------
|
||
# 4. 错误示例：并行节点覆盖同一个无 reducer 字段
|
||
# -----------------------------------------------------------------------------
|
||
|
||
|
||
class BadState(TypedDict):
|
||
value: str # 没有 reducer，默认策略是覆盖
|
||
|
||
|
||
def left_branch(state: BadState) -> dict:
|
||
return {"value": "来自左分支"}
|
||
|
||
|
||
def right_branch(state: BadState) -> dict:
|
||
return {"value": "来自右分支"}
|
||
|
||
|
||
bad_builder = StateGraph(BadState)
|
||
bad_builder.add_node("left", left_branch)
|
||
bad_builder.add_node("right", right_branch)
|
||
bad_builder.add_edge(START, "left")
|
||
bad_builder.add_edge(START, "right")
|
||
bad_builder.add_edge("left", END)
|
||
bad_builder.add_edge("right", END)
|
||
bad_graph = bad_builder.compile()
|
||
|
||
|
||
# -----------------------------------------------------------------------------
|
||
# 5. 演示函数
|
||
# -----------------------------------------------------------------------------
|
||
|
||
|
||
def demo_good_graph() -> None:
|
||
print("=== 正确示例：并行 + Reducer + 汇合 ===")
|
||
|
||
initial_state: ResearchState = {
|
||
"topic": "LangGraph",
|
||
"results": [],
|
||
"logs": [],
|
||
"summary": "",
|
||
}
|
||
|
||
started = time.perf_counter()
|
||
result = good_graph.invoke(initial_state)
|
||
elapsed = time.perf_counter() - started
|
||
|
||
print(f"\n耗时：{elapsed:.2f} 秒")
|
||
print("两个搜索节点各等待 1 秒；总耗时约 1 秒，说明它们并行执行。")
|
||
|
||
print("\nresults（operator.add，保留重复项）：")
|
||
for item in result["results"]:
|
||
print(f"- {item}")
|
||
|
||
print("\nlogs（merge_unique，重复日志只保留一次）：")
|
||
for item in result["logs"]:
|
||
print(f"- {item}")
|
||
|
||
print("\nsummary（默认覆盖，由 summarize 单独写入）：")
|
||
print(result["summary"])
|
||
|
||
|
||
def demo_stream_updates() -> None:
|
||
print("\n=== 使用 updates 观察并行 super-step ===")
|
||
|
||
initial_state: ResearchState = {
|
||
"topic": "LangGraph",
|
||
"results": [],
|
||
"logs": [],
|
||
"summary": "",
|
||
}
|
||
|
||
for part in good_graph.stream(
|
||
initial_state,
|
||
stream_mode="updates",
|
||
version="v2",
|
||
):
|
||
if part["type"] != "updates":
|
||
continue
|
||
|
||
node_names = ", ".join(part["data"].keys())
|
||
print(f"节点更新：{node_names}")
|
||
|
||
|
||
def demo_conflict() -> None:
|
||
print("\n=== 错误示例：并行更新冲突 ===")
|
||
try:
|
||
bad_graph.invoke({"value": "初始值"})
|
||
except InvalidUpdateError as exc:
|
||
print(f"捕获到：{type(exc).__name__}")
|
||
print("原因：left 和 right 在同一个 super-step 中都写入 value，")
|
||
print("而 value 没有 reducer，LangGraph 不知道应该保留哪一个值。")
|
||
print(f"原始错误：{exc}")
|
||
|
||
|
||
def main() -> None:
|
||
demo_good_graph()
|
||
demo_stream_updates()
|
||
demo_conflict()
|
||
|
||
print("\n=== 学习结论 ===")
|
||
print("1. 多条普通出边会让目标节点在同一个 super-step 中并行运行。")
|
||
print("2. 并行分支汇合时，下游节点等待本轮所有上游分支完成。")
|
||
print("3. 无 reducer 的字段采用覆盖策略，同一步只能接收一个更新。")
|
||
print("4. operator.add 可以连接列表，但不会去重。")
|
||
print("5. 自定义 reducer 可以实现去重、排序或其他合并规则。")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|