[拆解LangChain执行引擎] PregelProtocol——定义了"LangChain执行体"最小功能集

JaydenAI
0 阅读11分钟

Pregel是对PregelProtocol协议的实现，后者的引入标志着 LangGraph 从一个单一的库进化为了一个可插拔的图计算框架。图可以视为“LangGraph 执行体”，而PregelProtocol定义了它必须具备的最小功能集合。我们从这协议的成员定义来看看这个功能集合包含哪些操作。

1. 配置绑定

通过前面的内容我们会发现RunnableConfig这个对象几乎时无所不在，我们在调用Pregel对象的时候可以将它作为参数，用来提供用于控制其执行行为（比如迭代限制，并发控制等）的配置。执行引擎还将它作为容器用来下流流程传递一些组件和信号，所以前面的演示实例才可以在Node处理函数中从注入的RunnableConfig中提取像Runtime、PregelScratchpad、Checkpoint命名空间和静态上下文这样对象和信息。对于单纯Pregel的Node（不包括StateGraph的Node），RunnableConfig使唯一可以注入到处理函数中的参数，所以除了输入参数，其他所需的信息只能从它里面提取。

with_config方法赋予了这个 “执行体”与配置绑定的能力。除了提供RunnableConfig对象，我们还可以利用关键词参数提供待绑定的配置。由于RunnableConfig本质上就是一个TypedDict对象，提供的关键字参数组成的键值对可以直接转换成RunnableConfig对象。with_config方法会将两者合并，生成一个新的RunnableConfig对象绑定到执行体上。

class PregelProtocol(Runnable[InputT, Any], Generic[StateT, ContextT, InputT, OutputT]):
    @abstractmethod
    def with_config(
        self, config: RunnableConfig | None = None, **kwargs: Any
    ) -> Self: ...

2. 可视化呈现

PregelProtocol是LangGraph对 “图” 的抽象，这里的图是 “图论” 的概念，但是若真能将它的结构呈现在一张 “图片” 中，这无疑是非常有意义的。毕竟代码仅仅是面向程序员的语言，比不上图片，不但直观，还没有受众限制。LangGraph专门定义了如下这个Graph类型来表示面向 “可视化呈现” 的图。

一个Graph对象标识的图依然由Node和Edge构成。它的每个Node都有一个唯一标识，我们可以调用next_id方法为下一个待添加的Node生成此标识。我们不仅可以调用add_node、remove_node和add_edge这样的方法以添加/移除Node和Edge来构建图，还可以调用extend方法将另一个Graph的所有Node和Edge添加进来。

@dataclass
class Graph:
    nodes: dict[str, Node] = field(default_factory=dict)
    edges: list[Edge] = field(default_factory=list)
    
    def next_id(self) -> str
    def add_node(
        self,
        data: type[BaseModel] | RunnableType | None,
        id: str | None = None,
        *,
        metadata: dict[str, Any] | None = None,
    ) -> Node
    def remove_node(self, node: Node) -> None
    def add_edge(
        self,
        source: Node,
        target: Node,
        data: Stringifiable | None = None,
        conditional: bool = False,  # noqa: FBT001,FBT002
    ) -> Edge
    def extend(
        self, graph: Graph, *, prefix: str = ""
    ) -> tuple[Node | None, Node | None]:

    def reid(self) -> Graph:
    def first_node(self) -> Node | None
    def last_node(self) -> Node | None
    def trim_first_node(self) -> None
    def trim_last_node(self) -> None

    def to_json(self, *, with_schemas: bool = False) -> dict[str, list[dict[str, Any]]]
    def draw_ascii(self) -> str:
    def print_ascii(self) -> None:
    @overload
    def draw_png(
        self,
        output_file_path: str,
        fontname: str | None = None,
        labels: LabelsDict | None = None,
    ) -> None: ...
    @overload
    def draw_png(
        self,
        output_file_path: None,
        fontname: str | None = None,
        labels: LabelsDict | None = None,
    ) -> bytes: ...
    def draw_png(
        self,
        output_file_path: str | None = None,
        fontname: str | None = None,
        labels: LabelsDict | None = None,
    ) -> bytes | None
    def draw_mermaid(
        self,
        *,
        with_styles: bool = True,
        curve_style: CurveStyle = CurveStyle.LINEAR,
        node_colors: NodeStyles | None = None,
        wrap_label_n_words: int = 9,
        frontmatter_config: dict[str, Any] | None = None,
    ) -> str
    def draw_mermaid_png(
        self,
        *,
        curve_style: CurveStyle = CurveStyle.LINEAR,
        node_colors: NodeStyles | None = None,
        wrap_label_n_words: int = 9,
        output_file_path: str | None = None,
        draw_method: MermaidDrawMethod = MermaidDrawMethod.API,
        background_color: str = "white",
        padding: int = 10,
        max_retries: int = 1,
        retry_delay: float = 1.0,
        frontmatter_config: dict[str, Any] | None = None,
        base_url: str | None = None,
        proxies: dict[str, str] | None = None,
    ) -> bytes

调用reid方法可以返回一个新的Graph对象，它尽量保留途中可读性的元素，但是Node的ID会重新生成。Graph的first_node和last_node方法返回第一个和最后一个Node。如果我们希望删除第一个只有单一输出Edge或者最后一个只有单一输入Edge的Node，可以调用trim_first_node或者trim_last_node方法。

构建好的Graph可以采用不同的呈现方式。Graph提供了五个“绘图”方法，其中draw_ascii和print_ascii采用ascii码字符的呈现方式，前者返回具体的ascii码字符串，后者则直接在终端将图绘制出来，这种方法不依赖其他的绘图相关的包。draw_mermaid和draw_mermaid_png采用Mermaid图表的呈现方式，Mermaid 是一种基于文本的流程图定义语言，广泛支持于 GitHub、Notion 和各种编辑器中。draw_mermaid返回图标文本，而draw_mermaid_png则直接将图表进一步渲染成PNG图片。Graph对象也可以通过调用draw_png方法渲染成PNG图片，该方法最终会Graphviz（一个开源的图可视化软件）来布局和渲染图片。

再回到PregelProtocol类型的定义上，它定义了如下所示的get_graph/aget_graph方法，它们的返回类型DrawableGraph正是上述Graph类型的别名。该方法除了可以传入RunnableConfig对象作为可选配置外，还具有一个名为xray的参数。xray（X光）参数决定了你在查看图结构时，到底能看多深。它专门用于处理子图的展开显示。如果设置为False（默认值），图将以 “黑盒” 模式显式，如果你的图中包含子图，它只会显示为一个单一的节点。你看不见子图内部的任何节点、边或逻辑。反之将会采用 “全展开” 模式，它会像 X 光一样穿透所有层级，将所有嵌套子图内部的节点和连线全部平铺出来。

from langchain_core.runnables.graph import Graph as DrawableGraph
class PregelProtocol(Runnable[InputT, Any], Generic[StateT, ContextT, InputT, OutputT]):
    @abstractmethod
    def get_graph(
        self,
        config: RunnableConfig | None = None,
        *,
        xray: int | bool = False,
    ) -> DrawableGraph: ...

    @abstractmethod
    async def aget_graph(
        self,
        config: RunnableConfig | None = None,
        *,
        xray: int | bool = False,
    ) -> DrawableGraph: ...

在第一个演示实例中，我们创建了一个作为“笑话生成器”的Agent，现在我们将它简化，看看由它生成的Graph如何将图的结构以可视化的形式呈现出来。如下面的代码片段所示，我们利用StateGraph作为Builder，构建了一张由两个Node组成的图，它们和Start和End之间有四条边。

from langgraph.graph import StateGraph, START, END
from langgraph.pregel.protocol import PregelProtocol
from PIL import Image as PILImage
import io
from langgraph.checkpoint.memory import MemorySaver

def generate_joke(state):
    pass

def regenerate_joke(state):
    pass

builder = (
    StateGraph(dict)
    .add_node("generate_joke", generate_joke)
    .add_node("regenerate_joke", regenerate_joke)
)

builder.add_edge(START, "generate_joke")
builder.add_edge("regenerate_joke", END)
builder.add_conditional_edges(
    "generate_joke", lambda _: "bad", {"good": END, "bad": "regenerate_joke"}
)

app: PregelProtocol = builder.compile(MemorySaver())
graph = app.get_graph()
graph.print_ascii()

bytes = graph.draw_mermaid_png()
PILImage.open(io.BytesIO(bytes)).show()

在将StateGraph编译成Pregel对象后，我们调用其get_graph方法得到对应的Graph对象。我们以两种形式呈现其结构，前者通过调用print_ascii方法以ASCII字符的形式输出图结构，后者调用draw_mermaid_png方法生成一张PNG图片。下图左右两部分分别展现了两种呈现方式的效果。

3. 持久化

为了支持“中断/恢复”的执行方式，同时为“时间旅行”提供支持，图必须利用持久化的方式将执行过程的重要时刻的状态保存下来。LangGraph采用基于Checkpoint的持久化形式，对于指定的每个任务，不论是执行成功针对Channel的写入意图，还是抛出异常、人为中断或者Resume Value的提供，都会以Pending Write的形式被记录下来；当超步成功完成，针对Channel的写入被成功应用，这些Pending Write被丢弃，换来一个Checkpoint来描述当前的状态。

作为“LangGraph 执行体”的抽象，PregelProtocol定义了get_state/aget_state方法用于读取在某个Superstep由Checkpoint（对于最后一个未完成的Superstep，还包括Pending Write）构建的状态快照，该快照体现为一个StateSnapshot对象。get_state_history/aget_state_history返回由这些快照谱写的一段历史。

class PregelProtocol(Runnable[InputT, Any], Generic[StateT, ContextT, InputT, OutputT]):
    @abstractmethod
    def get_state(
        self, config: RunnableConfig, *, subgraphs: bool = False
    ) -> StateSnapshot: ...

    @abstractmethod
    async def aget_state(
        self, config: RunnableConfig, *, subgraphs: bool = False
    ) -> StateSnapshot: ...

    @abstractmethod
    def get_state_history(
        self,
        config: RunnableConfig,
        *,
        filter: dict[str, Any] | None = None,
        before: RunnableConfig | None = None,
        limit: int | None = None,
    ) -> Iterator[StateSnapshot]: ...

    @abstractmethod
    def aget_state_history(
        self,
        config: RunnableConfig,
        *,
        filter: dict[str, Any] | None = None,
        before: RunnableConfig | None = None,
        limit: int | None = None,
    ) -> AsyncIterator[StateSnapshot]: ...

    @abstractmethod
    def bulk_update_state(
        self,
        config: RunnableConfig,
        updates: Sequence[Sequence[StateUpdate]],
    ) -> RunnableConfig: ...

    @abstractmethod
    async def abulk_update_state(
        self,
        config: RunnableConfig,
        updates: Sequence[Sequence[StateUpdate]],
    ) -> RunnableConfig: ...

    @abstractmethod
    def update_state(
        self,
        config: RunnableConfig,
        values: dict[str, Any] | Any | None,
        as_node: str | None = None,
    ) -> RunnableConfig: ...

    @abstractmethod
    async def aupdate_state(
        self,
        config: RunnableConfig,
        values: dict[str, Any] | Any | None,
        as_node: str | None = None,
    ) -> RunnableConfig: ...

持久化存储的Checkpoint不仅使我们可以回顾历史，还可以提供“时间旅行”，使我们可以从某个历史时刻重新执行后面的流程。不仅如此，PregelProtocol还提供了update_state /bulk_update_state/abulk_update_state可以直接修改状态。但是它们并非“篡改历史”，只是基于某个在某个历史时刻开启了另一段“平行宇宙”而已。持久化使LangGraph.Pregel作为核心和部分，我们将在后续部分对它进行专门的介绍。

4. 两种调用方式

PregelProtocol的invoke/ainvoke和stream/astream方法体现了针对 “LangGraph 执行体” 两种调用方式。前者采用简单的请求/回复消息交换模式，客户端需要等整个流程结束之后采用得到结果。如果整个处理流程比较复杂，或者涉及一些耗时的操作，过长的等待会带来糟糕的体验。后者采用流式处理使客户端可以实施得到处理的中间结果或者感知到处理的进度。我们将在后续部分对流式处理进行单独介绍。

class PregelProtocol(Runnable[InputT, Any], Generic[StateT, ContextT, InputT, OutputT]):
    @abstractmethod
    def stream(
        self,
        input: InputT | Command | None,
        config: RunnableConfig | None = None,
        *,
        context: ContextT | None = None,
        stream_mode: StreamMode | list[StreamMode] | None = None,
        interrupt_before: All | Sequence[str] | None = None,
        interrupt_after: All | Sequence[str] | None = None,
        subgraphs: bool = False,
    ) -> Iterator[dict[str, Any] | Any]: ...

    @abstractmethod
    def astream(
        self,
        input: InputT | Command | None,
        config: RunnableConfig | None = None,
        *,
        context: ContextT | None = None,
        stream_mode: StreamMode | list[StreamMode] | None = None,
        interrupt_before: All | Sequence[str] | None = None,
        interrupt_after: All | Sequence[str] | None = None,
        subgraphs: bool = False,
    ) -> AsyncIterator[dict[str, Any] | Any]: ...

    @abstractmethod
    def invoke(
        self,
        input: InputT | Command | None,
        config: RunnableConfig | None = None,
        *,
        context: ContextT | None = None,
        interrupt_before: All | Sequence[str] | None = None,
        interrupt_after: All | Sequence[str] | None = None,
    ) -> dict[str, Any] | Any: ...

    @abstractmethod
    async def ainvoke(
        self,
        input: InputT | Command | None,
        config: RunnableConfig | None = None,
        *,
        context: ContextT | None = None,
        interrupt_before: All | Sequence[str] | None = None,
        interrupt_after: All | Sequence[str] | None = None,
    ) -> dict[str, Any] | Any: ...

执行体支持中断/恢复（interrupt/resume）的方式执行，所以在中断时需要将当时的状态以 “Checkpoint（Checkpoint）” 的形式保存下来，恢复执行的时候利用它们将当时的 “执行线程” 复原。Checkpointing的机制也使 “时间旅行” 成为可能，我们可以从任一Checkpoint开始执行。也正是因为此持久化机制的存在，我们可以提取某一个Superstep的状态，还可以查看整个执行历史，这两个功能分别对应PregelProtocol的get_state/aget_state和get_state_history/aget_state_history方法。具体的状态以StateSnapshot对象描述的快照表示。

执行体应该具有将执行结果作为新的状态进行保存的能力，所以PregelProtocol定义了update_state/aupdate_state和bulk_update_state/abulk_update_state方法，前者保存单一状态更新，后者对多个状态更新进行批量执行。单一状态更新通过如下这个名为StateUpdate的命名元组表示，我们不仅可以利用values字段得到以字典形式表示的状态值，还可以通过as_node和task_id字段的得到实施更新的Node和具体任务标识。

class StateUpdate(NamedTuple):
    values: dict[str, Any] | None
    as_node: str | None = None
    task_id: str | None = None

执行体支持两种基本的操作，一种采用单纯的请求/响应消息交换模式，另一种以流的形式实时返回数据，它们分别对应invoke/ainvoke和stream/astream方法。

5. 嵌套结构

我们一直在强调图的“嵌套”结构，这种结构也可以从Pregel、PregelNode和PregelProtocol在三个类型的定义。一个Pregel是PregelProtocol的实现、作为其节点的PregelNode对象可以由一个或者多个PregelProtocol组成，对于表示 “子图” 的subgraphs字段，并且该字段返回一个PregelProtocol对象的序列。Pregel的subgraphs方法返回的子图就来源于组成它的Node。

class PregelNode:
    subgraphs	: Sequence[PregelProtocol]

class Pregel(
    PregelProtocol[StateT, ContextT, InputT, OutputT],
    Generic[StateT, ContextT, InputT, OutputT]):
    def get_subgraphs(
        self, *, namespace: str | None = None, recurse: bool = False
    ) -> Iterator[tuple[str, PregelProtocol]]

    async def aget_subgraphs(
        self, *, namespace: str | None = None, recurse: bool = False
    ) -> AsyncIterator[tuple[str, PregelProtocol]]

PregelNode的subgraphs字段提供了 “子图” 的静态注册，其实任何一个Pregel对象都可以在无需注册前提下被另一个Pregel的Node调用，而且反映当前执行上下文的一些执行配置会通过上下文变量（ContenxtVars） “流向” 作为子图的Pregel对象。前面我们演示子图调用涉及的Checkpoint命名空间的例子已经充分体现了这一点。但是这种显式的静态声明对于图的静态图分析与可视化有着积极的作用。

avatar
文章
阅读
粉丝