如何创建强类型LangChain Runnable并实现多步骤数据传递？

修正后的完整代码 #

from typing import Optional
from pydantic import BaseModel, Field
from langchain_core.runnables import RunnableSerializable, RunnableConfig, RunnablePassthrough

# 假设以下为业务中已定义的类
class Analysis(BaseModel):
    map: dict

class Element(BaseModel):
    pass

class InterestingElement(Element):
    pass

class ElementSelectionContext(BaseModel):
    element_name: str = Field(frozen=True)
    objective: str = Field(frozen=True)
    page: Optional[Analysis] = None

class Page(RunnableSerializable[ElementSelectionContext, Analysis]):
    def invoke(self, input: ElementSelectionContext, config: RunnableConfig | None = None) -> Analysis:
        # 保留原文件读取逻辑，这里简化为返回示例数据
        # current_dir = os.path.dirname(os.path.abspath(__file__))
        # file_path = os.path.join(current_dir, 'sample.json')
        # with open(file_path, 'r', encoding="utf-8") as file:
        #     data = json.load(file)
        #     return Analysis(**data)
        return Analysis(map={"key1": InterestingElement(), "key2": InterestingElement()})

class ExactMatch(RunnableSerializable[ElementSelectionContext, Optional[InterestingElement]]):
    def invoke(self, input: ElementSelectionContext, config: RunnableConfig | None = None) -> Optional[InterestingElement]:
        if not input.page or not input.page.map:
            return None
        return next(iter(input.page.map.values()))

class ElementSelectionPipeline(RunnableSerializable[ElementSelectionContext, Optional[InterestingElement]]):
    def invoke(self, input: ElementSelectionContext, config: RunnableConfig | None = None) -> Optional[InterestingElement]:
        page_step = Page()
        exact_match_step = ExactMatch()
        
        pipeline = (
            # 执行page_step，将结果合并到原始上下文的字典结构中
            RunnablePassthrough.assign(page=page_step)
            # 将字典转回ElementSelectionContext强类型实例
            | lambda ctx_dict: ElementSelectionContext(**ctx_dict)
            # 传递更新后的上下文给ExactMatch
            | exact_match_step
        )
        
        return pipeline.invoke(input)

# 使用示例
pipeline = ElementSelectionPipeline()
response = pipeline.invoke(ElementSelectionContext(element_name="joba", objective="joba"))
print(response)

关键修改说明 #

• 修复字典输入问题：
  原代码中{"page": RunnablePassthrough()}会将输出转为纯字典，导致ExactMatch无法接收强类型的ElementSelectionContext。改用RunnablePassthrough.assign(page=page_step)可以保留原始上下文的键值对，同时新增page字段，再通过lambda将字典转回ElementSelectionContext实例，确保类型正确。

• 实现上下文更新传递：
  由于ElementSelectionContext的element_name和objective字段是冻结的，不能直接修改。通过RunnablePassthrough.assign自动合并原始上下文和page_step的结果，再用ElementSelectionContext(**ctx_dict)创建新实例，自然完成page字段的更新，同时保留冻结字段的值。

• 适配输入类型：
  修改ExactMatch的泛型定义和invoke方法的输入参数类型为ElementSelectionContext，与管道传递的输入匹配，同时添加空值判断避免运行时异常。

内容的提问来源于stack exchange，提问作者JobaDiniz