Доступ к переменным LangChain LCEL из предыдущих этапов цепочки

Я пытаюсь сесть на поезд LCEL и Langserve, но у меня возникли небольшие проблемы с пониманием некоторой «магии», связанной с доступом к переменным, установленным в словаре конвейера.

Переменные кажутся разрешимыми из шаблонов подсказок. Я хотел бы получить эти значения в пользовательских функциях и т. д., но мне не ясно, как получить к ним прямой доступ. Возьмем следующий надуманный пример, целью которого является возврат исходного документа, а также ответа в ответе:

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(answer):
    sources = [] # TODO lookup source documents (key: 'context')
    return ChatResponse(answer=answer, sources=sources)

retrieval_chain = (
        {"context": retriever, "question": RunnablePassthrough()}
        | prompt
        | llm
        | StrOutputParser()
        | RunnableLambda(format_response)
)
app = FastAPI()
add_routes(app, retrieval_chain, path = "/chat", input_type=str, output_type=ChatResponse)

В format_response я оставил TODO для поиска исходных документов. Я хотел бы получить исходные документы из ключа context конвейера. Как мне получить доступ к этому ключу, который был установлен на первом этапе цепочки?