langchainpy-langchain

How to pass input parameters through RunnableSequences in LangChain v0.2?

I am struggeling with basic chaining and passing input parameters through RunnableSequences in LangChain v0.2.

I have two chains: code_chain and test_chain.

enter image description here enter image description here

Now I want to chain them together.

enter image description here

This is my current code:

code_prompt = PromptTemplate.from_template("Write a very short {language} function that will {task}");
code_chain = code_prompt | llm | {"code": StrOutputParser()};

test_prompt = PromptTemplate.from_template("Write a test for the following {language} code:\n{code}");
test_chain = test_prompt | llm | {"test": StrOutputParser()};

chain = code_chain | test_chain;

result = chain.invoke({
    "language": "python",
    "task": "reverse a string",
});

Because code_chain does not retain the language parameter as an output, it is missing in the test_chain:

KeyError: "Input to PromptTemplate is missing variables {'language'}. Expected: ['code', 'language'] Received: ['code']"

How do I pass the language input of the first chain to the language input of the second?

Solution

  • It turns out my question is a duplicate of How to include the inputs of the first chain to the second chain in LangChain's SequentialChain?.

    To fully control the flow of variables through the pipeline we can use RunnableParallel objects:

    chain = RunnableParallel({
    "language": RunnablePick("language"),
    "code": code_chain | RunnablePick("code"),
}) | RunnableParallel({
    "language": RunnablePick("language"),
    "code": RunnablePick("code"),
    "test": test_chain | RunnablePick("test"),
});

    Now we have all variables available in our output: language, code, and test.