构建 SQL Agent

构建一个可以回答 SQL 数据库相关问题的 Agent。

本教程将教你如何使用 LangChain Agent 构建一个能够回答 SQL 数据库相关问题的 Agent。

Agent 的工作流程：

1. 从数据库获取可用表和 schema
2. 判断哪些表与问题相关
3. 获取相关表的 schema
4. 基于问题生成 SQL 查询
5. 双重检查查询的常见错误
6. 执行查询并返回结果
7. 纠正错误直到查询成功
8. 基于结果生成回复

安装

pip install langchain langgraph

配置数据库

本教程使用 SQLite 的 Chinook 数据库：

import requests, pathlib
import sqlite3
from langchain.tools import tool

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

con = sqlite3.connect("Chinook.db")
cursor = con.cursor()
cursor.execute("SELECT name FROM sqlite_master WHERE type='table';")
tables = [row[0] for row in cursor.fetchall() if not row[0].startswith("sqlite_")]
print(f"Available tables: {tables}")
con.close()

创建数据库工具

@tool
def sql_db_list_tables() -> str:
    """列出数据库中的所有表"""
    con = sqlite3.connect("Chinook.db")
    try:
        cursor = con.cursor()
        cursor.execute("SELECT name FROM sqlite_master WHERE type='table';")
        tables = [row[0] for row in cursor.fetchall() if not row[0].startswith("sqlite_")]
        return ", ".join(tables)
    finally:
        con.close()

@tool
def sql_db_schema(table_names: str) -> str:
    """获取表的 schema 和示例行"""
    con = sqlite3.connect("Chinook.db")
    try:
        cursor = con.cursor()
        cursor.execute("SELECT name FROM sqlite_master WHERE type='table';")
        valid_tables = {row[0] for row in cursor.fetchall() if not row[0].startswith("sqlite_")}
        results = []
        for table in table_names.split(","):
            table = table.strip()
            if table not in valid_tables:
                results.append(f"Error: table {table!r} not found")
                continue
            cursor.execute("SELECT sql FROM sqlite_master WHERE type='table' AND name=?;", (table,))
            schema = cursor.fetchone()
            if schema:
                results.append(schema[0])
        return "\n\n".join(results)
    finally:
        con.close()

@tool
def sql_db_query(query: str) -> str:
    """执行 SQL 查询"""
    con = sqlite3.connect("Chinook.db")
    try:
        cursor = con.cursor()
        cursor.execute(query)
        return str(cursor.fetchall())
    finally:
        con.close()

创建 Agent

from langchain.agents import create_agent

agent = create_agent(
    model="openai:gpt-4o",
    tools=[sql_db_list_tables, sql_db_schema, sql_db_query],
    system_prompt="You are a SQL agent. Use the tools to answer questions about the database.",
)

运行

result = agent.invoke({
    "messages": [{"role": "user", "content": "How many artists are in the database?"}]
})
print(result["messages"][-1].content)

下一步

• 语义搜索 — 构建语义搜索引擎
• RAG Agent — 构建 RAG Agent
• Agents 智能体 — 深入了解 Agent