如何在LangChain多步骤SQL查询链中自定义输入输出变量名？

核心问题拆解 #

1. 内置create_sql_query_chain默认输出键为query，需要重命名为unvalidated_query供后续验证步骤使用
2. 验证阶段的Prompt无法直接访问中间生成的自定义变量，需要手动传递上下文

解决方案代码实现 #

from langchain_core.runnables import RunnablePassthrough, RunnableLambda
from langchain_core.output_parsers import StrOutputParser

# 1. 生成未验证SQL：将默认输出"query"重命名为"unvalidated_query"
generate_unvalidated_sql = RunnablePassthrough.assign(
    unvalidated_query=create_sql_query_chain
)

# 2. 验证链：将需要的变量传递给Prompt，确保能识别自定义的unvalidated_query
validated_query_chain = (
    RunnableLambda(lambda x: {
        "user_question": x["question"],
        "db_context": x["db_context"],
        "unvalidated_query": x["unvalidated_query"]
    })
    | validation_prompt 
    | llm 
    | StrOutputParser()
)

# 3. 执行查询链：接收验证后的SQL并执行（替换为你实际的执行逻辑）
execute_query_chain = (
    RunnableLambda(lambda x: {"query": x["validated_query"]})
    | your_execute_query_function  # 例如SQLDatabase的run方法或自定义执行链
)

# 4. 最终工作流：链式组合所有步骤，确保变量按依赖顺序传递
final_chain = (
    generate_unvalidated_sql
    .assign(validated_query=validated_query_chain)
    .assign(result=execute_query_chain)
    .assign(
        answer=RunnableLambda(lambda x: {
            "user_question": x["question"],
            "query_result": x["result"]
        }) | answer_back_to_natural_language_prompt | llm | StrOutputParser()
    )
    # 筛选最终输出的字段
    | RunnableLambda(lambda x: {
        "unvalidated_query": x["unvalidated_query"],
        "validated_query": x["validated_query"],
        "result": x["result"],
        "answer": x["answer"]
    })
)

关键实现说明 #

• 重命名默认输出：通过RunnablePassthrough.assign将create_sql_query_chain的默认query输出映射为自定义的unvalidated_query，让后续步骤能直接引用这个变量。
• 传递变量到Prompt：用RunnableLambda将上下文里的所有必要变量（用户问题、数据库上下文、未验证SQL）重新组织成Prompt期望的输入格式，确保Prompt模板中的占位符（如{unvalidated_query}）能正确取值。
• 链式依赖控制：使用.assign()的链式调用，保证步骤按顺序执行——先生成未验证SQL，再验证，接着执行查询，最后转换为自然语言。同一assign内的步骤是并行执行的，因此依赖前序结果的步骤必须放在后续的.assign()中。
• 自定义输出筛选：最后通过RunnableLambda筛选需要返回的字段，避免返回冗余的上下文变量。

注意事项 #

• 确保你的validation_prompt模板中使用的变量名和传递的键完全一致，例如模板中要包含{unvalidated_query}才能正确获取未验证的SQL。
• 如果执行查询的函数需要特定的输入键（比如默认的query），用RunnableLambda将自定义的validated_query映射为函数所需的键名即可。

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