[SOLVED] Aggregating output from langchain LCEL elements

Aggregating output from langchain LCEL elements

I have two chains, one that generates a document and one that creates a short document resume. I want to chain them, using the output from the first on inside the other one. But I want to get both outputs in the result.

Before LCEL, I could do it using LLMChain's output_key parameter. With LCEL, there seems to be a RunnablePassthrough class, but I don't seem to get how to use it to aggregate the output. Code example:

generate_document_chain = generate_document_prompt | llm | StrOutputParser()
resume_document_chain = resume_document_prompt | llm | StrOutputParser()

aggregated_chain = generate_document_chain | resume_document_chain 
content = aggregated_chain.invoke({"topic": topic})

Solution

Perhaps the following is what you want. It feeds the output of the first chain into second chain as input.

from langchain_core.runnables import RunnablePassthrough

aggregated_chain = generate_document_chain | {
    "first_chain_output": RunnablePassthrough(),
    "second_chain_output": resume_document_chain
}
content = aggregated_chain.invoke({"topic": topic})

Then the output will be a dictionary with keys: "first_chain_output" and "second_chain_output".

You can also use RunnablePassthrough.assign. Unlike the case above, the key of generate_document_chain output should match the input variable name of the second chain. Below, the input variable of the second chain is assumed to be "input" (btw, the input variable of the first chain is "topic").

aggregated_chain = (
    {"input": generate_document_chain} 
    | RunnablePassthrough.assign(second_chain_output=resume_document_chain)
)

The output of this chain will be a dict with keys: "input" and "second_chain_output".