[拆解LangChain执行引擎] 静态上下文在Pregel中的应用

JaydenAI
0 阅读4分钟

在 Pregel 模型中，静态上下文是一个专门设计的依赖注入容器。它的出现是为了解决在复杂的图计算中，如何优雅地处理“不属于图状态，但Node运行又必须依赖的外部环境信息”这一痛点。这些数据具有一个共同的性质，那就是在整个运行生命周期内只读且固定，比如：

  • 身份信息：当前发起请求的user_id、org_id。
  • 外部客户端：已实例化的db_connection、redis_client、vector_store。
  • 策略约束：当前任务的safety_level或budget_limit。

静态上下文是 Pregel 运行时提供的一个类型安全的环境变量容器。它将执行环境（Context）与业务轨迹（State）物理隔离，使得大型 Agent 系统的架构更加模块化，也更容易在复杂的生产环境下进行测试和调试。

不同于以往字典形式的配置，静态上下文采用强类型 Schema 定义方法。由于其静态只读的特性，它在整个生命周期内保持一致性。静态上下文具有单次运行锁定机制，这保证Pregel对象一旦被调用，上下文对象在所有Node、所有 Superstep中引用的是同一个内存地址。它的非持久化特性进一步确保它不会被写入Checkpoint，所以当Pregel因为错误停止并从断点恢复时，我们必须重新提供一个相同的上下文对象。综上所示，静态上下文作为非序列化的、运行时的旁路注入而存在。

静态上下文在Pregel被作为Runtime的一部分来传递的。如下所示的Runtime类的泛型参数ConextT指的就是静态上下文数据类型。除了返回该上下文的context字段，Runtime还具有额外三个字段分别返回用于长期存储的store字段（返回一个BaseStore对象）、实现“custom”流模式的stream_writer字段（返回一个StreamWriter对象），以及提供当前会话上一个返回值的previous字段。

@dataclass(**_DC_KWARGS)
class Runtime(Generic[ContextT]):
    context: ContextT = field(default=None
    store: BaseStore | None = field(default=None)
    stream_writer: StreamWriter = field(default=_no_op_stream_writer)
    previous: Any = field(default=None)

Pregel节点的处理函数读取静态上下文比较繁琐，以为除了承载输入的参数（一般是一个字典），我们只能额外定义一个RunnableConfig类型的参数，意味着基本上出原始输入外的其他任务信息都得从这个RunnableConfig配置中提取。RunnableConfig是一个字典，所以我们要提取所需数据的前提是得预先知道对用得Key。这样设计也能理解，因为LangGraph.Prege在整个LangChain宇宙中作为执行引擎而存在，它相当于LangChain体系的内核。Pregel提供的API本就不是针对Agent应用开发者，对开发者友好不是Pregel得设计目标，保持这个内核足够简洁更重要。

RunnableConfig对象会贯穿整个Pregel引擎的执行，上游流程利用这个它像下游传递所需的组件和控制信息，传递的信息大都被至于configurable子节点下。如果对应的Key以__pregel_作为前缀，表示该条目其实是由Pregel内部使用的。Runtime对应的Key为__pregel_runtime。

如下这个例子演示了如何声明、指定和读取静态上下文。我们定义一个承载基本用户信息的UserInfo数据类型作为静态上下文的Schema。作为Pregel唯一的Node，其处理函数提供了一个RunnableConfig类型的参数，我们从中提供作为运行时的Runtime对象，进而得到作为静态上下文的UserInfo对象。

from langchain_core.runnables import RunnableConfig
from langgraph.pregel import Pregel, NodeBuilder
from typing import Any, Literal
from langgraph.channels import LastValue
from langgraph.runtime import Runtime
from dataclasses import dataclass

@dataclass
class UserInfo:
    id: str
    name: str
    gender: Literal["male", "female"]

def handle(args: dict[str, Any], config: RunnableConfig) -> str:
    runtime: Runtime = config["configurable"]["__pregel_runtime"]
    return runtime.context.__repr__()

node = (NodeBuilder()
    .subscribe_only("start")
    .write_to("output")
    .do(handle))
    
app = Pregel(
    nodes={"body": node},
    channels={"start": LastValue(None), "output": LastValue(str)},
    input_channels=["start"],
    output_channels=["output"],
    context_schema=UserInfo,
)

user = UserInfo(id="123", name="Alice", gender="female")
result = app.invoke(input={"start": None}, context=user)
assert result["output"] == user.__repr__()

在创建Pregel对象的时候，作为静态上下文的UserInfo类型直接以构造函数的context_schema参数进行声明。在调用其invoke方法的时候就通过context参数将指定的UserInfo对象作为静态上下文传递。静态上下文的设计初衷就是为了规避序列化的限制。它允许我们将复杂的、重量级的、带有外部依赖的对象的直接注入，而不会破坏 Pregel 模型对状态一致性和可持久化的要求。

avatar
文章
阅读
粉丝