自定义 SQL Agent

直接在 LangGraph 中实现 SQL Agent，获得最大灵活性。

本教程展示如何使用 LangGraph 原语构建一个 SQL Agent，可以回答数据库相关问题。

安装

pip install langchain langgraph langchain-openai

创建数据库工具

import sqlite3
from langchain.tools import tool

@tool
def list_tables() -> str:
    """列出数据库中所有表"""
    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()]
    con.close()
    return ", ".join(tables)

@tool
def get_schema(table_names: str) -> str:
    """获取表的 schema 信息"""
    con = sqlite3.connect("Chinook.db")
    cursor = con.cursor()
    results = []
    for table in table_names.split(","):
        table = table.strip()
        cursor.execute(
            "SELECT sql FROM sqlite_master WHERE type='table' AND name=?;",
            (table,)
        )
        row = cursor.fetchone()
        if row:
            results.append(row[0])
    con.close()
    return "\n\n".join(results)

@tool
def run_query(query: str) -> str:
    """执行 SQL 查询并返回结果"""
    con = sqlite3.connect("Chinook.db")
    cursor = con.cursor()
    cursor.execute(query)
    result = str(cursor.fetchall())
    con.close()
    return result

创建 Agent

from langchain.agents import create_agent

agent = create_agent(
    model="openai:gpt-4o",
    tools=[list_tables, get_schema, run_query],
    system_prompt="You are a SQL agent. Use the provided tools to answer questions.",
)

运行

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

下一步

• SQL Agent（LangChain） — 使用 LangChain Agent 构建 SQL 应用
• Agents 智能体 — 深入了解 Agent
• 语义搜索 — 构建语义搜索引擎