[拆解LangChain执行引擎-18]Agent状态是如何从通道读取的？

在Agent状态是如何被写入通道的？中，我们探讨了组成的Agent的节点在完成了自身操作后，如何将执行结果通过写入相应的通道，进而更新Agent的状态。本篇文章则讨论这个问题的另一面：节点如何通过读取通道得到Agent的状态的。

1. 利用ChannelRead读取状态

通过Agent状态是如何被写入通道的？，我们知道作为Agent节点的PregelNode对象利用ChannelWrite提交它对通道的写入意图，进而增量更新Agent的状态。ChannelRead与之相对，它被PregelNode用于读取相应的通道进而得到所需的状态成员，这个操作可以通过如下所示的静态方法do_read来完成。

class ChannelRead(RunnableCallable):
    @staticmethod
    def do_read(
        config: RunnableConfig,
        *,
        select: str | list[str],
        fresh: bool = False,
        mapper: Callable[[Any], Any] | None = None,
    ) -> Any

静态方法do_read具有如下的参数：

• config：当前的运行时配置；
• select：选择需要读取的通道，如果没有指定，将不会读取任何通道；
• fresh：如果设为True，它会强制获取当前状态下最新鲜的数据，而不是可能被缓存的旧值。默认值为False，以提高性能；
• mapper：对读取的内容（字典形式）作进一步转换，以生成最终返回的对象；

下面的程序演示了针对do_read方法的使用。我们创建了一个单节点的Pregel对象，并为它定义了5个通道，其中foo、bar、baz和select为输入通道，前三个通道用于提供常规的输入，select通道用于提供需要被节点读取的通道列表。节点订阅select通道，并将该通道的值作为参数调用do_read方法，作为运行时配置的config参数直接注入到节点函数。do_read方法返回的结果被写入result通道，该通道也是Pregel的唯一输出通道。

from langgraph.pregel import Pregel, NodeBuilder
from langgraph.pregel._read import ChannelRead
from langgraph.channels import LastValue
from typing import  Any
from langchain_core.runnables import RunnableConfig

def read(channels:list[str], config: RunnableConfig) -> Any:
    return ChannelRead.do_read(config, select=channels)

node = (NodeBuilder()
        .subscribe_only("select")
        .write_to("result")
        .do(read))

app = Pregel(
    nodes={"node": node},
    channels={
        "foo":LastValue(str),
        "bar":LastValue(str),
        "baz":LastValue(str),
        "select":LastValue(list),
        "result":LastValue(Any)
    },
    input_channels=["foo", "bar", "baz","select"],
    output_channels=["result"],
)

input = {"foo": "123", "bar": "456", "baz": "789"}

result = app.invoke(input={**input, "select": ["foo"]})
assert result["result"] == {"foo": "123"}

result = app.invoke(input={**input, "select": ["foo", "bar"]})
assert result["result"] == {"foo": "123", "bar": "456"}

result = app.invoke(input={**input, "select": ["foo","bar", "baz"]})    
assert result["result"] == {"foo": "123", "bar": "456", "baz": "789"}

result = app.invoke(input={**input, "select": []})
assert result["result"] == {}

Pregel对象创建之后，我们利用作为输入的字典为foo、bar和baz指定相同的值，但设置不同的select列表。我们将输入字典作为参数对Pregel发起多次调用，断言验证了返回值是否与select列表的一致性。

2. ChannelRead

静态方法do_read在ChannelRead中的定义非常简单，针对选择通道的读取并未真正实现在该方法中，一切都隐藏在它通过RunnableConfig中提取出来的这个可执行对象。

class ChannelRead(RunnableCallable):
    @staticmethod
    def do_read(
        config: RunnableConfig,
        *,
        select: str | list[str],
        fresh: bool = False,
        mapper: Callable[[Any], Any] | None = None,
    ) -> Any:
        try:
            read: READ_TYPE = config[CONF][CONFIG_KEY_READ]
        except KeyError:
            raise RuntimeError(
                "Not configured with a read function"
                "Make sure to call in the context of a Pregel process"
            )
        if mapper:
            return mapper(read(select, fresh))
        else:
            return read(select, fresh)

RunnableConfig在LangChain执行引擎中被当成数据总线在用，如果你对其底层实现稍有了解，你会发现它将很多东西往里面扔。我不想评论这种设计的优劣，但我个人是不太喜欢的。在Agent状态是如何被写入通道的？中，我们知道它在配置中为ChannelWrite的do_write方法存放了一个用于收集（channel_name, channel_value）二元组的函数，对应的Key为__pregel_send，通过常量CONFIG_KEY_SEND表示。这里也一样，它在配置存放了一个同于读取通道的函数，对应的Key为__pregel_read，通过常量CONFIG_KEY_READ表示。通过给出的实现，可以推断这个函数的签名为Callable[[list[str],bool],dict[str,Any]|Any]。

3. local_read

在Agent状态是如何被写入通道的？提到过，用于收集通道写入的函数是在创建节点对应执行任务的PregelExecutableTask对象时创建的，携带此函数的RunnableConfig也被绑定在这个PregelExecutableTask对象上。这个用于通道读取的函数的创建和应用方式与之类似。这个被置于RunnableConfig中的Callable[[list[str],bool],dict[str,Any]|Any]是针对如下这个local_read函数创建的偏函数。

def local_read(
    scratchpad: PregelScratchpad,
    channels: Mapping[str, BaseChannel],
    managed: ManagedValueMapping,
    task: WritesProtocol,
    select: list[str] | str,
    fresh: bool = False,
) -> dict[str, Any] | Any

ManagedValueMapping = dict[str, ManagedValueSpec]

local_read函数具有如下的参数：

• scratchpad：在推进Pregel执行的每一步（Superstep）对会创建的PregelScratchpad,用于存放当前超步数以及服务于中断恢复执行和SubAgent执行的计数器。所有ManagedValue的值都是根据它计算出来的；
• channels：所有通道与名称的映射；
• managed：所有ManagedValue的描述对象（ManagedValueSpec）与名称的映射；
• task：节点任务的协议，主要偏重于描述节点产生的状态更新；
• select：选择读取的通道列表；
• fresh：如果设为 True，它会强制获取当前状态下最新鲜的数据，而不是可能被缓存的旧值。默认值为False，以提高性能；

WritesProtocol协议定义如下，上面提到过的用于表示节点执行任务的PregelExecutableTask就实现了这个协议。它提供的四个属性分别表示任务的路径（SubAgent导致）、名称、通道写入（channel_name+channel_value二元组序列）和订阅通道。

class WritesProtocol(Protocol):
    @property
    def path(self) -> tuple[str | int | tuple, ...]: ...
    @property
    def name(self) -> str: ...
    @property
    def writes(self) -> Sequence[tuple[str, Any]]: ...
    @property
    def triggers(self) -> Sequence[str]: ...

如下所示的是local_read函数的完整实现：它利用select参数从channels和managed筛选出通道和ManagedValueSpec，然后调用read_channels将选择的通道和ManagedValue的值读出来。如果fresh参数被设置成True，它会对选择出来的通道先实施写操作来清除缓存。

def local_read(
    scratchpad: PregelScratchpad,
    channels: Mapping[str, BaseChannel],
    managed: ManagedValueMapping,
    task: WritesProtocol,
    select: list[str] | str,
    fresh: bool = False,
) -> dict[str, Any] | Any:
    updated: dict[str, list[Any]] = defaultdict(list)
    if isinstance(select, str):
        managed_keys = []
        for c, v in task.writes:
            if c == select:
                updated[c].append(v)
    else:
        managed_keys = [k for k in select if k in managed]
        select = [k for k in select if k not in managed]
        for c, v in task.writes:
            if c in select:
                updated[c].append(v)
    if fresh:
        # apply writes
        local_channels: dict[str, BaseChannel] = {}
        for k in channels:
            cc = channels[k].copy()
            cc.update(updated[k])
            local_channels[k] = cc
        # read fresh values
        values = read_channels(local_channels, select)
    else:
        values = read_channels(channels, select)
    if managed_keys:
        values.update({k: managed[k].get(scratchpad) for k in managed_keys})
    return values

在为节点创建任务的时候，它会提取节点的名称、订阅通道和当前的PregelScratchpad、所有通道和描述所有ManagedValue的ManagedValueSpec与其名称的映射字典，按照如下的方式调用创建local_read的偏函数。这个偏函数被置于绑定到节点任务的RunnableConfig配置中，供ChannelRead的静态方法do_read使用。代码中的那个writes变量就是那个用于收集通道写入意图的双端队列（deque）。

name = ...
writes = ...
triggers = ...
scratchpad = ...
channels = ...
managed = ...

partial(
    local_read,
    scratchpad,
    channels,
    managed,
    PregelTaskWrites(
        task_path[:3],
        name,
        writes,
        triggers,
    ),
),

class PregelTaskWrites(NamedTuple):
    path: tuple[str | int | tuple, ...]
    name: str
    writes: Sequence[tuple[str, Any]]
    triggers: Sequence[str]

顺便说说为什么叫local？在LangGraph的并行执行模型中，当一个节点执行完毕并返回结果时，这些结果属于局部更新，尚未合并到全局状态中。

4. read_channels

local_read函数用于读取通道的read_channels函数定义如下。没什么特别之处，最终就是调用BaseChannel的get方法，并在此基础上做了一些必要的异常处理而已。

def read_channels(
    channels: Mapping[str, BaseChannel],
    select: Sequence[str] | str,
    *,
    skip_empty: bool = True,
) -> dict[str, Any] | Any:
    if isinstance(select, str):
        return read_channel(channels, select)
    else:
        values: dict[str, Any] = {}
        for k in select:
            try:
                values[k] = read_channel(channels, k, catch=not skip_empty)
            except EmptyChannelError:
                pass
        return values

        def read_channel(
    channels: Mapping[str, BaseChannel],
    chan: str,
    *,
    catch: bool = True,
) -> Any:
    try:
        return channels[chan].get()
    except EmptyChannelError:
        if catch:
            return None
        else:
            raise

5. 作为Runnable的ChannelRead

对于ChannelRead，我们仅仅介绍了它的静态方法do_read，而它自身是一个继承自RunnableCallable的可执行对象，从如下的定义可以看出，最终用来读取通道的_read和_aread方法最终还是调用的静态方法do_read。

class ChannelRead(RunnableCallable):
    channel: str | list[str]
    fresh: bool = False
    mapper: Callable[[Any], Any] | None = None

    def __init__(
        self,
        channel: str | list[str],
        *,
        fresh: bool = False,
        mapper: Callable[[Any], Any] | None = None,
        tags: list[str] | None = None,
    ) -> None:
        super().__init__(
            func=self._read,
            afunc=self._aread,
            tags=tags,
            name=None,
            trace=False,
        )
        self.fresh = fresh
        self.mapper = mapper
        self.channel = channel

    def get_name(self, suffix: str | None = None, *, name: str | None = None) -> str:
        if name:
            pass
        elif isinstance(self.channel, str):
            name = f"ChannelRead<{self.channel}>"
        else:
            name = f"ChannelRead<{','.join(self.channel)}>"
        return super().get_name(suffix, name=name)

    def _read(self, _: Any, config: RunnableConfig) -> Any:
        return self.do_read(
            config, select=self.channel, fresh=self.fresh, mapper=self.mapper
        )

    async def _aread(self, _: Any, config: RunnableConfig) -> Any:
        return self.do_read(
            config, select=self.channel, fresh=self.fresh, mapper=self.mapper
        )