LangChain SQL Agent 实战指南:让AI帮你写SQL查询

前言

最近在做数据分析的时候,经常需要写各种复杂的SQL查询。有时候面对一个陌生的数据库,光是理解表结构就要花不少时间。后来发现LangChain提供了一个SQL Agent的功能,可以让大模型帮我们自动生成和执行SQL查询,真的省了不少事。今天就来详细分享一下怎么用。

这个Agent到底做了什么?

简单来说,SQL Agent就是一个能够理解自然语言问题,然后自动和数据库交互的智能助手。它的工作流程大概是这样的:

  1. 探索数据库结构 - 先看看数据库里有哪些表
  2. 筛选相关表 - 根据你的问题判断需要查哪些表
  3. 获取表结构 - 查看这些表的字段信息
  4. 生成SQL查询 - 根据问题和表结构写出SQL
  5. 检查语法 - 用LLM再检查一遍SQL有没有问题
  6. 执行查询 - 运行SQL获取结果
  7. 错误处理 - 如果出错了,会自动修正重试
  8. 返回答案 - 把查询结果整理成自然语言回复

听起来是不是很智能?不过这里要注意一个重要的安全问题:

警告:让AI自动生成并执行SQL是有风险的!一定要确保数据库连接的权限是最小化的,只给必要的读权限。虽然不能完全消除风险,但可以大大降低。

环境准备

安装依赖包

首先装好需要的Python包:

pip install langchain langgraph langchain-community

配置LangSmith(可选但推荐)

如果想看到Agent执行的详细过程,建议配置一下LangSmith。这东西就像是给Agent加了个调试器,能看到每一步它都做了什么:

export LANGSMITH_TRACING="true"
export LANGSMITH_API_KEY="你的API密钥"

选择合适的大模型

这个Agent需要用到支持工具调用(tool-calling)的模型。目前主流的几个都支持,下面列举几个常用的:

OpenAI(推荐)

import os
from langchain.chat_models import init_chat_model

os.environ["OPENAI_API_KEY"] = "sk-..."

model = init_chat_model("gpt-4.1")

或者直接用类初始化:

from langchain_openai import ChatOpenAI

os.environ["OPENAI_API_KEY"] = "sk-..."

model = ChatOpenAI(model="gpt-4.1")

Anthropic Claude

from langchain.chat_models import init_chat_model

os.environ["ANTHROPIC_API_KEY"] = "sk-..."

model = init_chat_model("claude-sonnet-4-5-20250929")

Google Gemini

from langchain.chat_models import init_chat_model

os.environ["GOOGLE_API_KEY"] = "..."

model = init_chat_model("google_genai:gemini-2.5-flash-lite")

Azure OpenAI

如果公司用的是Azure的服务:

import os
from langchain.chat_models import init_chat_model

os.environ["AZURE_OPENAI_API_KEY"] = "..."
os.environ["AZURE_OPENAI_ENDPOINT"] = "..."
os.environ["OPENAI_API_VERSION"] = "2025-03-01-preview"

model = init_chat_model(
    "azure_openai:gpt-4.1",
    azure_deployment=os.environ["AZURE_OPENAI_DEPLOYMENT_NAME"],
)

后面的例子我都用OpenAI来演示,其他模型用法类似。

准备数据库

这里用SQLite做演示,因为它轻量级、不需要额外安装数据库服务器。我们使用经典的Chinook示例数据库(模拟的一个数字音乐商店)。

下载数据库文件

LangChain官方把这个数据库放在了公开的GCS上,直接下载就行:

import requests
import pathlib

url = "https://storage.googleapis.com/benchmarks-artifacts/chinook/Chinook.db"
local_path = pathlib.Path("Chinook.db")

if local_path.exists():
    print(f"{local_path} 已存在,跳过下载")
else:
    response = requests.get(url)
    if response.status_code == 200:
        local_path.write_bytes(response.content)
        print(f"下载完成,保存到 {local_path}")
    else:
        print(f"下载失败,状态码: {response.status_code}")

连接数据库

使用LangChain提供的数据库封装工具连接:

from langchain_community.utilities import SQLDatabase

db = SQLDatabase.from_uri("sqlite:///Chinook.db")

# 看看数据库的基本信息
print(f"数据库类型: {db.dialect}")
print(f"可用的表: {db.get_usable_table_names()}")
print(f'示例数据: {db.run("SELECT * FROM Artist LIMIT 5;")}')

输出结果:

数据库类型: sqlite
可用的表: ['Album', 'Artist', 'Customer', 'Employee', 'Genre', 'Invoice', 'InvoiceLine', 'MediaType', 'Playlist', 'PlaylistTrack', 'Track']
示例数据: [(1, 'AC/DC'), (2, 'Accept'), (3, 'Aerosmith'), (4, 'Alanis Morissette'), (5, 'Alice In Chains')]

可以看到数据库里有11张表,包括专辑、艺术家、客户、发票等。

配置数据库工具

LangChain提供了一套现成的SQL数据库工具包,包含了几个关键工具:

from langchain_community.agent_toolkits import SQLDatabaseToolkit

toolkit = SQLDatabaseToolkit(db=db, llm=model)

tools = toolkit.get_tools()

# 看看都有哪些工具
for tool in tools:
    print(f"{tool.name}: {tool.description}\n")

这会输出4个工具:

  1. sql_db_query - 执行SQL查询的工具。如果查询有错误,会返回错误信息,然后Agent可以根据错误重新生成查询。

  2. sql_db_schema - 获取指定表的结构和示例数据。注意要先用list_tables确认表是否存在。

  3. sql_db_list_tables - 列出数据库中所有的表名。

  4. sql_db_query_checker - 在执行查询前检查SQL语句是否正确。这是个很有用的工具,可以减少错误。

创建SQL Agent

create_agent函数可以快速创建一个ReAct风格的Agent。这种Agent的特点是能够根据反馈不断调整策略。

编写系统提示词

先给Agent写一个详细的系统提示词,告诉它该怎么做:

system_prompt = """
你是一个专门与SQL数据库交互的智能助手。
当收到一个问题时,你需要创建一个语法正确的{dialect}查询来执行,
然后查看查询结果并返回答案。除非用户明确指定要查询的数据量,
否则请始终将查询结果限制在{top_k}条以内。

你可以按相关列对结果进行排序,以返回数据库中最有价值的示例。
不要查询表中的所有列,只查询与问题相关的列。

在执行查询之前,你必须仔细检查查询语句。如果执行查询时出错,
请重写查询并重试。

不要执行任何数据修改语句(INSERT、UPDATE、DELETE、DROP等)。

开始时,你应该始终先查看数据库中有哪些表,不要跳过这一步。

然后查询最相关的表的结构信息。
""".format(
    dialect=db.dialect,
    top_k=5,
)

这个提示词很关键,它定义了Agent的行为规范:

  • 限制返回结果数量(避免查询太多数据)
  • 只查询需要的列(提高效率)
  • 执行前检查SQL
  • 禁止修改数据库
  • 先看表结构再查询

初始化Agent

from langchain.agents import create_agent

agent = create_agent(
    model,
    tools,
    system_prompt=system_prompt,
)

就这么简单,Agent就创建好了!

运行Agent

现在来测试一下,问个实际的问题:

question = "哪个音乐类型的歌曲平均时长最长?"

for step in agent.stream(
    {"messages": [{"role": "user", "content": question}]},
    stream_mode="values",
):
    step["messages"][-1].pretty_print()

运行后你会看到Agent的完整思考过程:

================================ Human Message =================================
哪个音乐类型的歌曲平均时长最长?

================================== Ai Message ==================================
Tool Calls:
  sql_db_list_tables (call_BQsWg8P65apHc8BTJ1NPDvnM)
 Call ID: call_BQsWg8P65apHc8BTJ1NPDvnM
  Args:

================================= Tool Message =================================
Name: sql_db_list_tables

Album, Artist, Customer, Employee, Genre, Invoice, InvoiceLine, MediaType, Playlist, PlaylistTrack, Track

================================== Ai Message ==================================
Tool Calls:
  sql_db_schema (call_i89tjKECFSeERbuACYm4w0cU)
 Call ID: call_i89tjKECFSeERbuACYm4w0cU
  Args:
    table_names: Track, Genre

================================= Tool Message =================================
Name: sql_db_schema

CREATE TABLE "Genre" (
    "GenreId" INTEGER NOT NULL,
    "Name" NVARCHAR(120),
    PRIMARY KEY ("GenreId")
)
/*
3 rows from Genre table:
GenreId    Name
1    Rock
2    Jazz
3    Metal
*/

CREATE TABLE "Track" (
    "TrackId" INTEGER NOT NULL,
    "Name" NVARCHAR(200) NOT NULL,
    "AlbumId" INTEGER,
    "MediaTypeId" INTEGER NOT NULL,
    "GenreId" INTEGER,
    "Composer" NVARCHAR(220),
    "Milliseconds" INTEGER NOT NULL,
    "Bytes" INTEGER,
    "UnitPrice" NUMERIC(10, 2) NOT NULL,
    PRIMARY KEY ("TrackId"),
    FOREIGN KEY("GenreId") REFERENCES "Genre" ("GenreId"),
    ...
)

================================== Ai Message ==================================
Tool Calls:
  sql_db_query_checker (call_G64yYm6R6UauiVPCXJZMA49b)
 Call ID: call_G64yYm6R6UauiVPCXJZMA49b
  Args:
    query: SELECT Genre.Name, AVG(Track.Milliseconds) AS AverageLength
           FROM Track
           INNER JOIN Genre ON Track.GenreId = Genre.GenreId
           GROUP BY Genre.Name
           ORDER BY AverageLength DESC
           LIMIT 5;

================================= Tool Message =================================
Name: sql_db_query_checker

SELECT Genre.Name, AVG(Track.Milliseconds) AS AverageLength
FROM Track
INNER JOIN Genre ON Track.GenreId = Genre.GenreId
GROUP BY Genre.Name
ORDER BY AverageLength DESC
LIMIT 5;

================================== Ai Message ==================================
Tool Calls:
  sql_db_query (call_AnO3SrhD0ODJBxh6dHMwvHwZ)
 Call ID: call_AnO3SrhD0ODJBxh6dHMwvHwZ
  Args:
    query: SELECT Genre.Name, AVG(Track.Milliseconds) AS AverageLength
           FROM Track
           INNER JOIN Genre ON Track.GenreId = Genre.GenreId
           GROUP BY Genre.Name
           ORDER BY AverageLength DESC
           LIMIT 5;

================================= Tool Message =================================
Name: sql_db_query

[('Sci Fi & Fantasy', 2911783.0384615385),
 ('Science Fiction', 2625549.076923077),
 ('Drama', 2575283.78125),
 ('TV Shows', 2145041.0215053763),
 ('Comedy', 1585263.705882353)]

================================== Ai Message ==================================

平均来说,"Sci Fi & Fantasy"(科幻与奇幻)类型的歌曲时长最长,
平均约2,911,783毫秒。其次是"Science Fiction"、"Drama"、"TV Shows"和"Comedy"。

看到了吗?Agent自己完成了所有步骤:

  1. 先查看了所有表
  2. 确定需要Track和Genre两张表
  3. 获取了这两张表的结构
  4. 生成了SQL查询
  5. 用checker检查了SQL
  6. 执行查询并返回结果
  7. 把结果转换成了自然语言

整个过程完全自动化,而且逻辑清晰。如果你想看更详细的执行轨迹,可以在LangSmith的追踪页面查看。

添加人工审核(重要!)

虽然Agent很智能,但在生产环境中,让AI直接执行SQL还是有风险的。有时候它可能生成一个效率很低的查询,或者不小心查询了敏感数据。

LangChain提供了"人工介入"(Human-in-the-loop)功能,可以让Agent在执行关键操作前暂停,等待人工确认。

配置中断机制

from langchain.agents import create_agent
from langchain.agents.middleware import HumanInTheLoopMiddleware
from langgraph.checkpoint.memory import InMemorySaver

agent = create_agent(
    model,
    tools,
    system_prompt=system_prompt,
    middleware=[
        HumanInTheLoopMiddleware(
            interrupt_on={"sql_db_query": True},  # 在执行查询前暂停
            description_prefix="工具执行需要审批",
        ),
    ],
    checkpointer=InMemorySaver(),  # 需要添加checkpointer来支持暂停/恢复
)

这里我们配置了在执行sql_db_query工具前暂停。

运行带审核的查询

question = "哪个音乐类型的歌曲平均时长最长?"
config = {"configurable": {"thread_id": "1"}}  # 需要一个线程ID来追踪状态

for step in agent.stream(
    {"messages": [{"role": "user", "content": question}]},
    config,
    stream_mode="values",
):
    if "messages" in step:
        step["messages"][-1].pretty_print()
    elif "__interrupt__" in step:
        print("=== 等待审批 ===")
        interrupt = step["__interrupt__"][0]
        for request in interrupt.value["action_requests"]:
            print(request["description"])

运行后,Agent会在执行SQL前暂停:

...

=== 等待审批 ===
工具执行需要审批

Tool: sql_db_query
Args: {'query': 'SELECT g.Name AS Genre, AVG(t.Milliseconds) AS AvgTrackLength
                 FROM Track t
                 JOIN Genre g ON t.GenreId = g.GenreId
                 GROUP BY g.Name
                 ORDER BY AvgTrackLength DESC
                 LIMIT 1;'}

这时候你可以检查这个SQL,如果没问题就批准执行:

from langgraph.types import Command

for step in agent.stream(
    Command(resume={"decisions": [{"type": "approve"}]}),  # 批准执行
    config,
    stream_mode="values",
):
    if "messages" in step:
        step["messages"][-1].pretty_print()
    elif "__interrupt__" in step:
        print("=== 等待审批 ===")
        interrupt = step["__interrupt__"][0]
        for request in interrupt.value["action_requests"]:
            print(request["description"])

如果发现SQL有问题,也可以拒绝并让Agent重新生成:

Command(resume={"decisions": [{"type": "reject", "feedback": "这个查询可能会很慢,请优化一下"}]})

在LangSmith Studio中运行(可选)

如果想要一个更友好的交互界面,可以在LangSmith Studio中运行这个Agent。这样就有了一个聊天界面,可以连续提问,还能看到生成的SQL和执行结果。

配置步骤

首先安装额外的依赖:

pip install -U langgraph-cli[inmem]>=0.4.0

在项目目录创建langgraph.json

{
  "dependencies": ["."],
  "graphs": {
    "agent": "./sql_agent.py:agent"
  },
  "env": ".env"
}

创建sql_agent.py文件:

import pathlib
from langchain.agents import create_agent
from langchain.chat_models import init_chat_model
from langchain_community.agent_toolkits import SQLDatabaseToolkit
from langchain_community.utilities import SQLDatabase
import requests

# 初始化模型
model = init_chat_model("gpt-4.1")

# 下载数据库
url = "https://storage.googleapis.com/benchmarks-artifacts/chinook/Chinook.db"
local_path = pathlib.Path("Chinook.db")

if local_path.exists():
    print(f"{local_path} 已存在,跳过下载")
else:
    response = requests.get(url)
    if response.status_code == 200:
        local_path.write_bytes(response.content)
        print(f"文件已下载并保存为 {local_path}")
    else:
        print(f"下载失败。状态码: {response.status_code}")

db = SQLDatabase.from_uri("sqlite:///Chinook.db")

# 创建工具
toolkit = SQLDatabaseToolkit(db=db, llm=model)
tools = toolkit.get_tools()

# 系统提示词
system_prompt = """
你是一个专门与SQL数据库交互的智能助手。
当收到一个问题时,你需要创建一个语法正确的{dialect}查询来执行,
然后查看查询结果并返回答案。除非用户明确指定要查询的数据量,
否则请始终将查询结果限制在{top_k}条以内。

你可以按相关列对结果进行排序,以返回数据库中最有价值的示例。
不要查询表中的所有列,只查询与问题相关的列。

在执行查询之前,你必须仔细检查查询语句。如果执行查询时出错,
请重写查询并重试。

不要执行任何数据修改语句(INSERT、UPDATE、DELETE、DROP等)。

开始时,你应该始终先查看数据库中有哪些表,不要跳过这一步。

然后查询最相关的表的结构信息。
""".format(
    dialect=db.dialect,
    top_k=5,
)

# 创建agent
agent = create_agent(
    model,
    tools,
    system_prompt=system_prompt,
)

然后运行:

langgraph dev

就可以在浏览器中访问可视化界面了。

实践建议

安全性

  1. 最小权限原则:数据库用户只给SELECT权限,绝不给UPDATE、DELETE、DROP等权限
  2. 敏感数据:包含敏感信息的表(如用户密码、支付信息)最好不要暴露给Agent
  3. 添加审核:生产环境强烈建议加上人工审核机制
  4. 限制资源:设置查询超时时间、结果行数限制,防止资源耗尽

性能优化

  1. 索引优化:确保数据库表有合适的索引,否则Agent生成的查询可能很慢
  2. 结果限制:在system_prompt中明确限制返回行数(比如top_k=5)
  3. 缓存机制:如果经常问类似的问题,可以考虑加缓存

提示词调优

系统提示词很重要,可以根据你的数据库特点定制:

  • 如果表名或字段名有特殊命名规则,在提示词中说明
  • 如果有一些常用的查询模式,可以在提示词中给例子
  • 如果某些表有特殊含义,可以在提示词中解释

错误处理

Agent虽然会自动重试,但有时候还是会失败。常见问题:

  • 表名错误:Agent可能猜错表名,确保让它先执行list_tables
  • 字段类型不匹配:有时候Agent会搞混数字和字符串,需要在schema中说明清楚
  • 复杂JOIN:对于复杂的多表关联,Agent可能理解不准确,这时候可能需要人工介入

总结

LangChain的SQL Agent是个很实用的工具,特别适合:

  • 数据分析师:快速探索陌生数据库
  • 产品经理:不用写SQL也能查数据
  • 开发人员:快速生成查询原型

但是要注意:

  • 不要在生产数据库上直接用,最好用只读副本
  • 添加必要的安全措施和审核机制
  • 对于关键业务查询,还是建议人工编写SQL

如果你想更深入定制Agent的行为,可以直接用LangGraph的底层API来实现,灵活性会更高。

扩展阅读

希望这篇文章对你有帮助!如果有问题欢迎交流。

Logo

有“AI”的1024 = 2048,欢迎大家加入2048 AI社区

更多推荐