除了我们显式声明的用于存储业务数据或驱动信号的通道之外,Pregel自身也会维护一些系统通道,其中最重要的莫过于一个名为__pregel_tasks的通道。通过前面针对BSP的介绍,我们知道当Superstep进入同步屏障并应用所有更新后,引擎会根据节点针对通道的订阅情况和通道自身的更新状态生成下一步待执行的任务,其实待执行的任务不限于此。

1. 两种任务创建方式

我们将根据节点针对通道的订阅来驱动任务执行的模式称为Pull模式,与之相对的则是借助于__pregel_tasks这个通道实现的Push模式。具体来说,这是一个关闭累积模式的Topic类型的通道,它存储的Topic体现为具有如下定义的Send对象。当某个节点执行之后,可以像这个通道中写入一个Send来驱动某个节点在下一Superstep中执行。除了利用Send对象的node字段指定待执行的节点名称外,还可以利用arg字段提供输入参数。

class Send:   
    node: str
    arg: Any
    def __init__(self, /, node: str, arg: Any) -> None

由于关闭了累积模式,在Topic类型通道中写入的内容只会在下一个Superstep中生效,并且具有阅后即焚的特性。对于执行引擎来说,这个名为__pregel_tasks的通道存储的就是下一Superstep以Push模式驱动执行的任务列表,两者完美契合。

2. 确认__pregel_tasks通道的存在

__pregel_tasks通道的存在可以通过如下的演示实例来验证。如代码片段所示,在采用常规方式将Pregel对象创建出来后,我们根据通道名称从它的channels字段中将此通道提取出来。断言揭示了该通道自身的类型、存储的数据类型和累积模式开关。

from langgraph.channels import LastValue, Topic
from langgraph.pregel import Pregel, NodeBuilder
from langgraph.types import Send, Sequence

node = (
    NodeBuilder()
    .subscribe_only("input_channel")
    .do(lambda args: args)
    .write_to("output_channel")
)
app = Pregel(
    nodes={"node": node},
    channels={"input_channel": LastValue(str), "output_channel": LastValue(str)},
    input_channels=["input_channel"],
    output_channels=["output_channel"],
)

tasks: Topic[Send] = app.channels["__pregel_tasks"]
assert isinstance(tasks, Topic)
assert tasks.ValueType == Sequence[Send]
assert tasks.accumulate == False

3. 被保护起来的通道

虽然__pregel_tasks就是一个普通的Topic类型的通道,但是它并未开发对外部使用,Pregel把它保护得非常好。我们不能声明一个与之同名的通道,否则就会像如下的方式一样抛出一个ValueError,并提示Channel '__pregel_tasks' is reserved and cannot be used in the graph.

from langgraph.channels import Topic
from langgraph.pregel import Pregel, NodeBuilder
from langgraph.types import Send, Sequence

try:
    app = Pregel(
        nodes={"node": NodeBuilder().subscribe_only("__pregel_tasks")},
        channels={"__pregel_tasks": Topic[Sequence[Send]]},
        input_channels=["input_channel"],
        output_channels=["output_channel"],
    )
    assert False, "Expected an error due to reserved channel name"
except Exception as e:
    assert isinstance(e, ValueError)
    assert str(e) == "Channel '__pregel_tasks' is reserved and cannot be used in the graph."

我们也不能采用常规的方式将向其发送Send对象。比如在如下的演示程序中,节点foo试图向此通道发送一个驱动节点bar执行的Send对象,最终抛出了一个InvalidUpdateError异常,并提示Cannot write to the reserved channel TASKS。除此之外,由于Pregel在利用它将基于Push模式的任务创建出来后就会将其清空,所以我们也无法读取其中的任务。

from langgraph.channels import LastValue
from langgraph.pregel import Pregel, NodeBuilder
from langgraph.types import Send
from langgraph.errors import InvalidUpdateError

foo = (NodeBuilder()
    .subscribe_to("start",read= False)
    .do(lambda _: Send(node="bar", arg="foobar"))
    .write_to("__pregel_tasks"))
bar = (NodeBuilder()
    .do(lambda args:args)
    .write_to("output"))

app = Pregel(
    nodes={ "foo": foo, "bar": bar},
    channels={
        "start": LastValue(str),
        "output": LastValue(str),
    },
    input_channels=["start"],
    output_channels=["output"])
try: 
    app.invoke({"start": None})
    assert False, "Should have raised InvalidUpdateError"
except Exception as e:
    assert isinstance(e, InvalidUpdateError)
    assert str(e) == "Cannot write to the reserved channel TASKS"

4. 唯一的解决方案

我们能够想到的常规方法针对此通道的写入基本都绕不开引擎针对它的保护机制。我们在以Actor模型的角度来看Pregel中提到过,节点利用ChannelWriter对象实现针对通道的写入。我们可以将针对通道的写入意图封装成ChannelWriteTupleEntry,并以此来创建ChannelWriter,这应该是唯一能够欺骗引擎写入验证的手段。如代码片段所示,率先执行的节点foo会返回一个驱动节点bar指定的Send对象,为了将它写入__pregel_tasks,我们创建了一个ChannelWriter,针对该通道的写入定义在ChannelWriteTupleEntry对象中,具体体现在调用构造函数指定的mapper参数上,它提供一个映射将节点的执行结果转成通道名称和值的映射关系。

from langgraph.pregel import Pregel, NodeBuilder
from langgraph.channels import LastValue
from langgraph.pregel._read import PregelNode
from langgraph.pregel._write import ChannelWrite, ChannelWriteTupleEntry
from langgraph.types import Send

foo: PregelNode = (NodeBuilder()
    .subscribe_to("foo")
    .do(lambda _: Send(node="bar", arg="foo"))
).build()
entry = ChannelWriteTupleEntry(mapper= lambda args: [("__pregel_tasks", args)])
foo.writers.append(ChannelWrite(writes=[entry]))

bar = (NodeBuilder()
    .do(lambda args: f"bar is triggered by {args}.")
    .write_to("output"))

app = Pregel(
    nodes={"foo": foo, "bar": bar},
    channels={
        "foo": LastValue(None),
        "output": LastValue(str),
    },
    input_channels=["foo"],
    output_channels=["output"],
)

result = app.invoke(input={"foo": None})
assert result == {"output": "bar is triggered by foo."}
Logo

Agent 垂直技术社区,欢迎活跃、内容共建。

更多推荐