How can I merge intermediate chain step outputs into a flat JSON output using LangChain?

Below is a short explanation, along with an example, on how you can use LangChain’s dictionary‐merging functionality to update your data like a dict.update() in the pipeline.

When you run several sub‑chains with RunnableParallel, each branch returns its result and gets nested under its key. That’s why you saw a highly nested output. To “flatten” your combined information (i.e. merge the input with intermediate steps), you need a mechanism to update the input dictionary with the outputs of later steps. This is exactly what RunnablePassthrough.assign() offers: it “merges” computed outputs into the input dict so that you end up with one flat dictionary containing the original input and all the intermediary outputs.

For example, suppose you have these two operations (the first chain extracts names and the second chain validates them). Instead of running them in parallel (which nests the results), you can update the dictionary as the chain flows by using assign. Here’s an illustrative snippet:

---

PYTHONCopied
from langchain_core.runnables import RunnableLambda, RunnablePassthrough, RunnableParallel
from operator import itemgetter

# Imagine these are your operations.
# The first chain extracts names from an input string.
first_chain = (
    first_prompt  # your PromptTemplate that extracts names
    | llm
    | SimpleJsonOutputParser()
)

# The second chain validates names against a second list.
second_chain = (
    second_prompt  # your PromptTemplate that compares lists
    | llm
    | SimpleJsonOutputParser()
)

# Start with a dictionary carrying the original inputs:
base = {"first_value": first_value, "possible_values": possible_values}

# First, run the extraction step and update (merge) the input dictionary.
chain1 = RunnablePassthrough.assign(first_prompt_output=first_chain)

# Then, run the validation step and merge its output as well.
chain2 = RunnablePassthrough.assign(second_prompt_output=second_chain)

# Chain them sequentially so that each step updates the dictionary.
# (This is equivalent to performing dict.update() as the data flows.)
final_chain = chain1 | chain2

# When you invoke the chain, the output would be a merged dictionary:
# {
#   "first_value": "Dave, John, carrot",
#   "possible_values": [...],
#   "first_prompt_output": ...,
#   "second_prompt_output": ...
# }
output = final_chain.invoke(base)
print(output)

---

In this example:

• The first assignment using RunnablePassthrough.assign() runs your first_chain and “updates” the input dict with its result under the key first_prompt_output.
• Then, the second assignment runs your second_chain and updates the dict with second_prompt_output.
• The final resulting dictionary contains all intermediary steps as well as the original input.

Key Points:

• RunnablePassthrough.assign() is used to merge new keys (and their computed values) into the input dictionary.
• When you use RunnableParallel, the separate branches build a nested dictionary because each branch’s result is stored under its branch key.
• If you prefer a flat, merged output (like dict.update()), then it’s best to use assign sequentially so that each new result is merged at the top level.

This approach keeps your data flow clear and ensures the final output is easier to read, while retaining all your intermediate steps along with the input.