如何修复LangChain串联链中代码与测试输出重复的问题？

问题根源 #

1. chain_code输出解析逻辑错误：你试图将LLM生成的代码内容同时解析为language和code字段，但LLM的输出仅包含生成的代码，language是初始输入的变量，无需从LLM输出中提取。
2. chain_test输出解析逻辑错误：你错误地将LLM生成的测试代码同时解析为test和code字段，导致原生成的代码被测试代码覆盖，最终出现code与test结果一致的问题。

修复后的代码 #

from langchain_openai import OpenAI
from langchain_core.prompts import PromptTemplate
from langchain_core.runnables import RunnablePassthrough
from langchain_core.output_parsers import StrOutputParser

llm = OpenAI()

code_prompt = PromptTemplate(
    input_variables=["task", "language"],
    template="Write a very short {language} function that will {task}."
)
test_prompt = PromptTemplate(
    input_variables=["language", "code"],
    template="Write a unit test for the following {language} code:\n{code}"
)

# 修复chain_code：保留输入的language，仅将LLM输出解析为code
chain_code = RunnablePassthrough.assign(
    code=code_prompt | llm | StrOutputParser()
)

# 修复chain_test：保留chain_code传递的code，将LLM输出解析为test
chain_test = RunnablePassthrough.assign(
    test=test_prompt | llm | StrOutputParser()
)

sequence = chain_code | chain_test

result = sequence.invoke({"task": "return a list of numbers", "language":"python"})

print(">>>>>>> GENERATED CODE:")
print(result["code"])

print(">>>>>>> GENERATED TEST:")
print(result["test"])

关键修复点说明 #

• 使用RunnablePassthrough.assign()方法：该方法会自动保留所有输入变量（比如初始的language、chain_code生成的code），同时新增指定的输出字段（code或test），避免了字段被错误覆盖的问题。
• 拆分解析逻辑：LLM的输出仅对应各自的目标字段（代码对应code，测试对应test），不再重复解析无关字段。

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