如何在LangChain LCEL链的自定义函数中访问前置步骤变量？

问题根源是原链经过StrOutputParser()后仅保留了回答文本，前面步骤的context（源文档）未传递到format_response函数中。要解决这个问题，需调整链的结构，让context和生成的回答一同传递到自定义函数里。

修改后的完整代码 #

from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnablePassthrough, RunnableLambda
from langchain_openai import ChatOpenAI, OpenAIEmbeddings
from langchain_community.vectorstores import FAISS
from langchain_core.documents import Document
from pydantic import BaseModel, List
from fastapi import FastAPI
from langserve import add_routes

class ChatResponse(BaseModel):
    answer: str
    sources: List[Document]

store = FAISS.from_texts(
    ["harrison worked at kensho"], embedding=OpenAIEmbeddings()
)
retriever = store.as_retriever()
template = """Answer the question based only on the following context:
{context}

Question: {question}
"""
prompt = ChatPromptTemplate.from_template(template)
llm = ChatOpenAI()

def format_response(data):
    # 直接从传入的字典中获取context对应的源文档
    sources = data["context"]
    return ChatResponse(answer=data["answer"], sources=sources)

retrieval_chain = (
    # 第一步：获取检索到的context和用户输入的question
    {"context": retriever, "question": RunnablePassthrough()}
    # 第二步：并行执行，生成answer并保留原始context
    | RunnablePassthrough.assign(
        answer=prompt | llm | StrOutputParser()
    )
    # 第三步：将包含context和answer的完整字典传给format_response
    | RunnableLambda(format_response)
)

app = FastAPI()
add_routes(app, retrieval_chain, path="/chat", input_type=str, output_type=ChatResponse)

关键说明 #

• RunnablePassthrough.assign()会把原始输入（包含context和question）与新生成的answer合并为一个字典，确保后续函数能访问所有需要的数据。
• 调整后的format_response函数接收字典参数，直接从中提取context（源文档列表）和answer，无需额外查询操作。

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