LangChain SQL Agent：如何配置自定义提示与对话历史以支持多轮问询？

可以直接基于现有Agent和Toolkit实现需求，不需要额外方法，具体调整如下：

一、添加自定义提示规则 #

通过修改SQL Agent的PromptTemplate，注入自定义规则（比如存储问题指定用表X）。LangChain的SQLDatabaseToolkit提供了默认prompt，我们可以在此基础上扩展：

from langchain.prompts import PromptTemplate

# 获取toolkit的默认prompt
base_prompt = toolkit.get_default_prompt()

# 自定义规则，替换成你的实际业务要求
custom_instructions = """
额外规则：
1. 当用户询问存储相关问题时，必须优先使用`storage_table`（替换为你的表X名称）进行查询。
2. 查询前先确认表结构，避免使用不存在的字段或表。
"""

# 拼接新的prompt，保留原有必要变量
custom_prompt = PromptTemplate(
    input_variables=base_prompt.input_variables,
    template=base_prompt.template + "\n" + custom_instructions
)

二、修复对话上下文关联问题 #

当前代码的ConversationBufferMemory已配置，但需确保Agent的prompt包含chat_history变量，且memory的key与prompt变量名一致。修改后的完整代码如下：

from langchain.memory import ConversationBufferMemory
from langchain.sql_database import SQLDatabase
from langchain.llms import OpenAI  # 替换为你实际使用的LLM

# 初始化数据库和LLM（替换为你的实际配置）
db = SQLDatabase.from_uri("your_db_uri")
llm = OpenAI(temperature=0)

toolkit = SQLDatabaseToolkit(db=db, llm=llm)
# 配置memory，指定key并设置返回消息格式
memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)

# 构建带上下文关联规则的自定义prompt
base_prompt = toolkit.get_default_prompt()
custom_instructions = """
额外规则：
1. 当用户询问存储相关问题时，必须优先使用`storage_table`进行查询。
2. 处理后续问题时，必须结合对话历史（chat_history），关联之前的提问内容作答。
"""
# 确保prompt包含chat_history变量
custom_prompt = PromptTemplate(
    input_variables=base_prompt.input_variables + ["chat_history"],
    template=base_prompt.template + "\n" + custom_instructions
)

# 创建Agent时指定自定义prompt和memory
agent_executor = create_sql_agent(
    llm=llm,
    toolkit=toolkit,
    verbose=True,
    memory=memory,
    prompt=custom_prompt
)

# 测试对话
agent_executor.run("How many products bought in Chicago")
agent_executor.run("这些产品是什么")  # 现在可关联上下文作答

关键说明 #

• 自定义提示：通过扩展默认prompt注入业务规则，LLM会遵循规则生成符合要求的SQL查询。
• 上下文保留：确保prompt包含chat_history变量，memory的memory_key与之匹配，同时设置return_messages=True让历史消息格式更清晰，便于Agent理解上下文关联关系。
• 若需控制历史消息长度，可改用ConversationBufferWindowMemory，避免冗余信息干扰LLM判断。

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