如何自定义Langchain SQL Agent限制，禁止查询表前三行？

针对你遇到的SQL Agent查询表前三行违反数据合规的问题，可通过以下几种方式处理：

1. 直接移除获取样本数据的工具 #

SQLDatabaseToolkit中包含一个用于获取表数据样本的工具（类名为SQLGetSampleDataTool），直接从工具列表中过滤掉它即可阻止相关查询操作：

from langchain.agents import AgentType, create_sql_agent
from langchain.sql_database import SQLDatabase
from langchain.agents.agent_toolkits.sql.toolkit import SQLDatabaseToolkit
from langchain.agents.agent_toolkits.sql.base import SQLGetSampleDataTool

# 初始化数据库与LLM实例
db = SQLDatabase.from_uri("你的数据库连接URI")
llm = 你的LLM模型实例

# 创建工具包并过滤违规工具
toolkit = SQLDatabaseToolkit(db=db, llm=llm)
filtered_tools = [tool for tool in toolkit.get_tools() if not isinstance(tool, SQLGetSampleDataTool)]

# 使用过滤后的工具创建合规的SQL Agent
agent = create_sql_agent(
    llm=llm,
    tools=filtered_tools,
    agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
    verbose=True
)

2. 替换违规工具为仅返回表结构的安全工具 #

如果担心移除工具影响Agent的推理逻辑，可自定义一个工具替代原样本查询工具，仅返回表结构信息而不涉及实际数据：

from langchain.tools import BaseTool
from langchain.agents.agent_toolkits.sql.base import SQLGetSampleDataTool
from typing import Optional

# 定义安全样本查询工具
class SafeSQLSampleTool(BaseTool):
    name = "sql_db_get_sample_data"
    description = "仅用于获取表的结构信息，无法返回实际数据样本"
    db: SQLDatabase

    def _run(self, table_name: str) -> str:
        # 仅返回表结构，不查询任何数据行
        table_schema = self.db.get_table_info(table_name)
        return f"表【{table_name}】的结构信息：\n{table_schema}\n\n合规提示：无法返回实际数据样本"

    async def _arun(self, table_name: str) -> str:
        return self._run(table_name)

# 替换原工具
toolkit = SQLDatabaseToolkit(db=db, llm=llm)
filtered_tools = []
for tool in toolkit.get_tools():
    if isinstance(tool, SQLGetSampleDataTool):
        filtered_tools.append(SafeSQLSampleTool(db=db))
    else:
        filtered_tools.append(tool)

# 创建Agent
agent = create_sql_agent(
    llm=llm,
    tools=filtered_tools,
    agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
    verbose=True
)

3. 修改数据库配置取消样本数据展示 #

初始化SQLDatabase时，可通过sample_rows_in_table_info参数设置为0，让表结构信息中不包含任何样本数据，从根源减少Agent触发样本查询的可能：

db = SQLDatabase.from_uri(
    "你的数据库连接URI",
    sample_rows_in_table_info=0  # 禁止在表结构中展示样本数据
)

建议优先使用第一种方案，直接移除违规工具，彻底避免查询敏感数据的风险；若需保留工具的语义提示，再结合第二种方案替换工具。

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