如何为LangChain的SQLDatabaseChain添加有效记忆并修复查询错误？

问题说明 #

需要实现带记忆功能的SQLDatabaseChain，完成两步关联查询：先查询指定域名网站的所有者，再查询该所有者的邮箱。但现有代码中，尽管ConversationBufferMemory已保存对话历史，第二次查询“Tell me his email”时，生成的SQL错误地查询了Westley Waters的邮箱，而非之前返回的Geo Mertz，导致结果错误。要求不使用Agent，通过简单链修复该问题。

原因分析 #

默认的SQLDatabaseChain提示词未包含对话历史的占位符，即使配置了ConversationBufferMemory，LLM生成SQL时也不会读取历史上下文，只能基于当前孤立的查询语句生成SQL，因此出现了错误的查询对象。

解决方案 #

自定义包含对话历史的提示词，让LLM在生成SQL前参考之前的对话内容，明确当前查询的上下文关联对象；同时确保SQLDatabaseChain正确加载这个自定义提示词，实现记忆与查询逻辑的绑定。

修改后的代码 #

import os
from langchain import OpenAI, SQLDatabase, SQLDatabaseChain, PromptTemplate
from langchain.memory import ConversationBufferMemory

# 自定义提示词，加入对话历史占位符，明确要求LLM基于上下文生成SQL
template = """Given the following conversation history and a new user question, generate the appropriate SQL query to answer the question.
Make sure to use the conversation history to understand context if the question refers to previous information.

Conversation History:
{history}

New Question: {input}
SQL Query:"""

PROMPT = PromptTemplate(
    input_variables=["history", "input"],
    template=template
)

memory = ConversationBufferMemory(memory_key="history")
db = SQLDatabase.from_uri(os.getenv("DB_URI"))
llm = OpenAI(temperature=0, verbose=True)

# 使用自定义提示词初始化SQLDatabaseChain
db_chain = SQLDatabaseChain.from_llm(
    llm,
    db,
    verbose=True,
    memory=memory,
    prompt=PROMPT
)

# 执行查询
db_chain.run("Who is owner of the website with domain https://damon.name")
db_chain.run("Tell me his email")
print(memory.load_memory_variables({}))

预期运行结果 #

> Entering new  chain...
Who is owner of the website with domain https://damon.name
SQLQuery:SELECT first_name, last_name FROM owners JOIN websites ON owners.id = websites.owner_id WHERE domain = 'https://damon.name' LIMIT 5;
SQLResult: [('Geo', 'Mertz')]
Answer:Geo Mertz is the owner of the website with domain https://damon.name.
> Finished chain.
    
> Entering new  chain...
Tell me his email
SQLQuery:SELECT email FROM owners WHERE first_name = 'Geo' AND last_name = 'Mertz'
SQLResult: [('geo.mertz@example.com',)]
Answer:Geo Mertz's email is geo.mertz@example.com.
> Finished chain.
{'history': "Human: Who is owner of the website with domain https://damon.name\nAI: Geo Mertz is the owner of the website with domain https://damon.name.\nHuman: Tell me his email\nAI: Geo Mertz's email is geo.mertz@example.com."}

关键改动说明 #

• 新增自定义提示词模板，加入{history}占位符，明确告知LLM要利用对话历史理解当前查询的上下文。
• 初始化PromptTemplate时指定history和input为输入变量，确保记忆内容能被注入到提示词中。
• 初始化SQLDatabaseChain时传入自定义的prompt参数，替换默认提示词，实现上下文关联。

补充说明 #

如果LLM仍偶尔出现上下文理解偏差，可以进一步优化提示词（比如明确要求“必须使用对话历史中的所有者姓名生成SQL查询”），或确认temperature已设为0以保证输出确定性。若简单链无法满足复杂多轮上下文需求，再考虑使用Agent方案。

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