第六章：LangGraph Channel 与 Pregel 引擎精读

class BaseChannel(ABC, Generic[Value, Update, C]):
    """
    Value    — 从 Channel 中读取时得到的类型
    Update   — 向 Channel 写入时接受的类型
    C        — 序列化到 Checkpoint 时的类型
    """
    
    @abstractmethod
    def get(self) -> Value:
        """读取当前值。如果 Channel 为空，抛出 EmptyChannelError"""
        ...
    
    @abstractmethod
    def update(self, values: Sequence[Update]) -> bool:
        """接收一批写入值，返回是否有变化"""
        ...
    
    @abstractmethod
    def checkpoint(self) -> C:
        """导出当前状态用于持久化"""
        ...
    
    @abstractmethod
    def from_checkpoint(self, checkpoint: C) -> Self:
        """从 checkpoint 恢复"""
        ...

class LastValue(BaseChannel[Value, Value, Value]):
    """
    只保留最后一次写入的值。
    核心约束：每个超步(superstep)中，最多只能有一个节点写入。
    """
    
    def __init__(self, typ: Type[Value]) -> None:
        self.typ = typ
        self.value: Value | None = None
    
    def get(self) -> Value:
        if self.value is None:
            raise EmptyChannelError()
        return self.value
    
    def update(self, values: Sequence[Value]) -> bool:
        if len(values) == 0:
            return False
        if len(values) > 1:
            raise InvalidUpdateError(
                "LastValue channel 在同一个超步中收到了多个写入！"
                "这通常意味着多个节点同时写入了同一个字段。"
            )
        self.value = values[0]
        return True
    
    def checkpoint(self) -> Value:
        return self.value

class BinaryOperatorAggregate(BaseChannel[Value, Value, Value]):
    """
    使用二元运算符（reducer）合并所有写入值。
    允许多个节点在同一个超步中同时写入。
    """
    
    def __init__(self, typ: Type[Value], operator: Callable[[Value, Value], Value]) -> None:
        self.typ = typ
        self.operator = operator  # 如 operator.add
        self.value: Value | None = None
    
    def update(self, values: Sequence[Value]) -> bool:
        if len(values) == 0:
            return False
        
        new_value = values[0]
        for v in values[1:]:
            new_value = self.operator(new_value, v)
        
        if self.value is None:
            self.value = new_value
        else:
            self.value = self.operator(self.value, new_value)
        
        return True

import operator
from typing import Annotated

class AgentState(TypedDict):
    # messages 字段使用 add 合并 → BinaryOperatorAggregate Channel
    messages: Annotated[list[BaseMessage], operator.add]

# EphemeralValue: 每个超步后自动清除，不持久化
class EphemeralValue(BaseChannel[Value, Value, Value]):
    """读取后即清空。用于一次性信号传递。"""
    pass

# Topic: 发布-订阅模式，可以有多个读者
class Topic(BaseChannel[Sequence[Value], Value, list[Value]]):
    """积累所有写入值，读取时返回全部，然后清空。"""
    pass

# Context: 只读 Channel，由外部注入，节点不能写入
class Context(BaseChannel[Value, Value, Value]):
    """用于注入配置、API key 等不变信息。"""
    pass

from typing import Annotated, TypedDict
import operator

class MyState(TypedDict):
    # 规则 1: 普通类型 → LastValue Channel
    current_plan: str               # → LastValue[str]
    
    # 规则 2: Annotated + 函数 → BinaryOperatorAggregate Channel  
    messages: Annotated[list, operator.add]   # → BinaryOperatorAggregate[list]
    
    # 规则 3: Annotated + 自定义 reducer
    recent: Annotated[list, lambda old, new: (old + new)[-5:]]
    #                                        ↑ 只保留最近 5 条

# 伪代码：StateGraph 初始化
def _create_channels(state_type: type) -> dict[str, BaseChannel]:
    channels = {}
    for field_name, field_type in get_type_hints(state_type).items():
        if is_annotated(field_type):
            # 提取 reducer 函数
            base_type, reducer = extract_annotated(field_type)
            channels[field_name] = BinaryOperatorAggregate(base_type, reducer)
        else:
            channels[field_name] = LastValue(field_type)
    return channels

class BadState(TypedDict):
    # 陷阱 1: 忘记 Annotated，多节点并行写入会报错！
    results: list  # → LastValue[list]，不是 BinaryOperatorAggregate！
    
    # 正确写法：
    results: Annotated[list, operator.add]  # → 支持多节点并行写入

# 陷阱 2: reducer 函数的返回类型必须与字段类型匹配
class AlsoBadState(TypedDict):
    count: Annotated[int, operator.add]  # ← 这其实是合法的！
    # int + int = int，operator.add 对 int 也适用
    # 但语义是"累加"而非"覆盖"——确保这是你想要的

# 简化版 tick() 逻辑
def tick(self) -> list[str]:
    """返回当前超步应该执行的节点列表"""
    
    ready_nodes = []
    
    for node_name, node in self.nodes.items():
        # 获取该节点订阅的 Channel 列表
        subscribed_channels = node.triggers
        
        # 获取该节点上次执行时看到的 Channel 版本
        last_seen_versions = self.checkpoint.versions_seen.get(node_name, {})
        
        # 判断：是否有任何订阅的 Channel 在上次执行后被更新了？
        should_run = False
        for ch_name in subscribed_channels:
            current_version = self.checkpoint.channel_versions.get(ch_name, 0)
            seen_version = last_seen_versions.get(ch_name, 0)
            
            if current_version > seen_version:
                should_run = True
                break
        
        if should_run:
            ready_nodes.append(node_name)
    
    return ready_nodes

# 简化版 after_tick() 逻辑
def after_tick(self, node_outputs: dict[str, dict]) -> None:
    """处理所有节点的执行结果"""
    
    # 1. Apply writes: 将节点输出写入 Channel
    for node_name, output in node_outputs.items():
        for channel_name, value in output.items():
            channel = self.channels[channel_name]
            
            # 调用 Channel 的 update 方法
            # LastValue: 直接覆盖（但如果多个节点写同一个 LastValue 会报错）
            # BinaryOperatorAggregate: 用 reducer 合并
            changed = channel.update([value])
            
            if changed:
                # 更新 Channel 版本号
                self.checkpoint.channel_versions[channel_name] += 1
        
        # 记录该节点看到的 Channel 版本
        for ch_name in self.nodes[node_name].triggers:
            self.checkpoint.versions_seen[node_name][ch_name] = \
                self.checkpoint.channel_versions[ch_name]
    
    # 2. 保存 Checkpoint
    if self.checkpointer:
        self.checkpointer.put(self.checkpoint)

@dataclass
class Checkpoint:
    # 唯一标识
    id: str                          # checkpoint UUID
    
    # Channel 的当前值
    channel_values: dict[str, Any]   
    # 例: {"messages": [...], "current_plan": "...", "draft": "..."}
    
    # Channel 的当前版本号
    channel_versions: dict[str, int]  
    # 例: {"messages": 5, "current_plan": 2, "draft": 3}
    # 含义：messages 被写入了 5 次，current_plan 被写入了 2 次...
    
    # 每个节点上次执行时看到的版本号
    versions_seen: dict[str, dict[str, int]]
    # 例: {
    #   "writer": {"messages": 3, "current_plan": 2},
    #   "reviewer": {"messages": 4, "draft": 2}
    # }
    # 含义：writer 上次执行时，messages 的版本是 3
    #       如果现在 messages 版本是 5（>3），writer 应该再次执行
    
    # 元信息
    ts: str                          # 时间戳
    parent_id: str | None            # 上一个 checkpoint 的 ID（形成链表）
    
    # 待处理的写入（用于 pending writes / interrupt 恢复）
    pending_writes: list[tuple[str, str, Any]]

# 图结构：
# start → agent → tools → agent → ... → end
# （典型 ReAct 循环）

# 初始状态
checkpoint = Checkpoint(
    channel_values={"messages": [HumanMessage("你好")]},
    channel_versions={"messages": 1},
    versions_seen={"agent": {}, "tools": {}},  # 都没执行过
)

# === 超步 1 ===
# tick(): agent 订阅 messages，版本 1 > 0（没见过） → agent 应该执行
# Execute: agent 读取 messages，调用 LLM，返回 AIMessage(tool_calls=[...])
# after_tick(): 
#   channel_values["messages"] = [..., AIMessage(tool_calls=[...])]
#   channel_versions["messages"] = 2
#   versions_seen["agent"]["messages"] = 2  ← agent 已经看到了版本 2

# === 超步 2 ===
# tick(): 
#   - agent: messages 版本 2 = seen 版本 2 → 不执行
#   - tools: 订阅 messages，版本 2 > 0（没见过） → tools 应该执行
# Execute: tools 执行工具调用，返回 ToolMessage(result="...")
# after_tick():
#   channel_values["messages"] = [..., ToolMessage(result="...")]
#   channel_versions["messages"] = 3
#   versions_seen["tools"]["messages"] = 3

# === 超步 3 ===
# tick():
#   - agent: messages 版本 3 > seen 版本 2 → 应该执行！
#   - tools: messages 版本 3 = seen 版本 3 → 不执行
# Execute: agent 读取包含工具结果的完整消息列表，决定是否需要更多工具调用
# ...

# 时间旅行示例
checkpointer = MemorySaver()
app = graph.compile(checkpointer=checkpointer)

# 执行
config = {"configurable": {"thread_id": "user-123"}}
result = app.invoke({"messages": [HumanMessage("帮我写代码")]}, config)

# 查看所有 checkpoint
for cp in checkpointer.list(config):
    print(f"超步 {cp.metadata['step']}: {cp.id}")

# 回到第 2 步重新执行
old_config = {"configurable": {"thread_id": "user-123", "checkpoint_id": "cp-step-2"}}
new_result = app.invoke({"messages": [HumanMessage("换个方案")]}, old_config)

from langgraph.graph import StateGraph, END
from langgraph.checkpoint.memory import MemorySaver

graph = StateGraph(AgentState)
graph.add_node("draft", draft_node)
graph.add_node("review", review_node)  # ← 我们想在这里暂停
graph.add_node("publish", publish_node)

graph.add_edge("draft", "review")
graph.add_edge("review", "publish")

# 编译时指定中断点
app = graph.compile(
    checkpointer=MemorySaver(),
    interrupt_before=["review"],  # review 节点执行前暂停
)

# 第一次执行
config = {"configurable": {"thread_id": "article-1"}}
result = app.invoke({"topic": "AI Agent 入门"}, config)
# 执行到 review 前暂停
# result 包含 draft 节点的输出

# 查看当前状态
state = app.get_state(config)
print(state.values["draft"])  # 查看草稿内容
print(state.next)             # ['review'] — 下一步要执行 review

# 方式 1: 直接恢复（不修改状态）
result = app.invoke(None, config)  # 传 None 表示"继续"

# 方式 2: 修改状态后恢复
app.update_state(config, {"draft": "人工修改后的草稿"})
result = app.invoke(None, config)

# 方式 3: 跳转到其他节点
app.update_state(config, {"draft": "..."}, as_node="review")
# 这告诉引擎："假装 review 已经执行了，直接去下一步"

# 伪代码：中断检查
def _should_interrupt(self, ready_nodes: list[str]) -> bool:
    for node in ready_nodes:
        if node in self.interrupt_before:
            # 保存当前 checkpoint
            self.checkpointer.put(self.checkpoint)
            # 抛出 GraphInterrupt
            raise GraphInterrupt(
                interrupts=[Interrupt(
                    value=f"等待确认后执行 {node}",
                    resumable=True,
                    ns=[node],
                )]
            )
    return False

def after_tick(self, outputs, executed_nodes):
    # ... 正常的 apply_writes 和 checkpoint ...
    
    for node in executed_nodes:
        if node in self.interrupt_after:
            raise GraphInterrupt(
                interrupts=[Interrupt(
                    value=outputs[node],  # 把节点输出作为中断值
                    resumable=True,
                )]
            )

from typing import Annotated, TypedDict
from langgraph.graph import StateGraph, END
import operator

class ResearchState(TypedDict):
    query: str
    sources: Annotated[list[str], operator.add]  # 多节点可同时写入
    summary: str                                  # 只能单节点写入
    quality_score: float

def search_web(state):
    return {"sources": [f"Web result for: {state['query']}"]}

def search_papers(state):
    return {"sources": [f"Paper result for: {state['query']}"]}

def summarize(state):
    all_sources = "\n".join(state["sources"])
    return {"summary": f"Summary of {len(state['sources'])} sources"}

def evaluate(state):
    score = 0.8 if len(state["sources"]) > 1 else 0.4
    return {"quality_score": score}

def route_after_eval(state):
    if state["quality_score"] >= 0.7:
        return "done"
    return "need_more"

# 构建图
graph = StateGraph(ResearchState)
graph.add_node("search_web", search_web)
graph.add_node("search_papers", search_papers)
graph.add_node("summarize", summarize)
graph.add_node("evaluate", evaluate)

graph.set_entry_point("search_web")
# search_web 和 search_papers 并行
graph.add_edge("search_web", "summarize")
graph.add_edge("search_papers", "summarize")
graph.set_conditional_entry_points(["search_web", "search_papers"])  # 并行入口

graph.add_edge("summarize", "evaluate")
graph.add_conditional_edges("evaluate", route_after_eval, {
    "done": END,
    "need_more": "search_web",
})

app = graph.compile()

# LangGraph: Checkpoint 包含完整的 Channel 状态 + 版本信息
# 支持 Time Travel（回到任意历史节点）
# 支持分支（从同一个 Checkpoint 产生不同的执行路径）

checkpoint = {
    "channel_values": {...},      # 所有 Channel 的当前值
    "channel_versions": {...},    # 全局版本号
    "versions_seen": {...},       # 每个节点的版本记录
    "pending_writes": [...],      # 未完成的写入
}

def sliding_window(max_size: int):
    """只保留最近 max_size 条消息的 reducer"""
    def reducer(old: list, new: list) -> list:
        combined = old + new
        return combined[-max_size:]  # 只保留最后 N 条
    return reducer

class ChatState(TypedDict):
    # 只保留最近 20 条消息，防止 context window 溢出
    messages: Annotated[list[BaseMessage], sliding_window(20)]
    
    # 只保留最近 5 次工具调用结果
    tool_results: Annotated[list[str], sliding_window(5)]

def deduplicate(old: list, new: list) -> list:
    """去重合并"""
    seen = set()
    result = []
    for item in old + new:
        key = item if isinstance(item, str) else id(item)
        if key not in seen:
            seen.add(key)
            result.append(item)
    return result

class SearchState(TypedDict):
    # 多个搜索节点可能返回重复结果，自动去重
    results: Annotated[list[str], deduplicate]

def accumulate_if_relevant(threshold: float):
    """只累积相关度超过阈值的结果"""
    def reducer(old: list, new: list) -> list:
        filtered_new = [item for item in new if item.get("relevance", 0) > threshold]
        return old + filtered_new
    return reducer

class RAGState(TypedDict):
    # 只保留相关度 > 0.7 的检索结果
    documents: Annotated[list[dict], accumulate_if_relevant(0.7)]

# 定义子图
research_subgraph = StateGraph(ResearchSubState)
research_subgraph.add_node("search", search_node)
research_subgraph.add_node("filter", filter_node)
research_subgraph.add_edge("search", "filter")
research_sub = research_subgraph.compile()

# 在主图中使用子图作为节点
main_graph = StateGraph(MainState)
main_graph.add_node("research", research_sub)  # 子图作为节点！
main_graph.add_node("write", write_node)
main_graph.add_edge("research", "write")

from langgraph.constants import Send

def route_to_workers(state):
    """动态决定启动几个 worker"""
    tasks = state["pending_tasks"]
    # 为每个任务动态创建一个 worker 执行
    return [Send("worker", {"task": task}) for task in tasks]

graph.add_conditional_edges("planner", route_to_workers)

# 流式输出：实时看到每个节点的结果
async for event in app.astream({"query": "什么是 Agent?"}, config):
    for node_name, output in event.items():
        print(f"[{node_name}] → {output}")

# 更细粒度：流式获取 LLM 的每个 token
async for event in app.astream_events(input, config, version="v2"):
    if event["event"] == "on_chat_model_stream":
        print(event["data"]["chunk"].content, end="")