构建 SQL 代理

CREATE TABLE "Album" (
	"AlbumId" INTEGER NOT NULL,
	"Title" NVARCHAR(160) NOT NULL,
	"ArtistId" INTEGER NOT NULL,
	PRIMARY KEY ("AlbumId"),
	FOREIGN KEY("ArtistId") REFERENCES "Artist" ("ArtistId")
)

/*
Album 表中的前 3 行数据：
AlbumId	Title	ArtistId
1	For Those About To Rock We Salute You	1
2	Balls to the Wall	2
3	Restless and Wild	2
*/


CREATE TABLE "Artist" (
	"ArtistId" INTEGER NOT NULL,
	"Name" NVARCHAR(120),
	PRIMARY KEY ("ArtistId")
)

/*
Artist 表中的前 3 行数据：
ArtistId	Name
1	AC/DC
2	Accept
3	Aerosmith
*/


CREATE TABLE "Customer" (
	"CustomerId" INTEGER NOT NULL,
	"FirstName" NVARCHAR(40) NOT NULL,
	"LastName" NVARCHAR(20) NOT NULL,
	"Company" NVARCHAR(80),
	"Address" NVARCHAR(70),
	"City" NVARCHAR(40),
	"State" NVARCHAR(40),
	"Country" NVARCHAR(40),
	"PostalCode" NVARCHAR(10),
	"Phone" NVARCHAR(24),
	"Fax" NVARCHAR(24),
	"Email" NVARCHAR(60) NOT NULL,
	"SupportRepId" INTEGER,
	PRIMARY KEY ("CustomerId"),
	FOREIGN KEY("SupportRepId") REFERENCES "Employee" ("EmployeeId")
)

/*
Customer 表中的前 3 行数据：
CustomerId	FirstName	LastName	Company	Address	City	State	Country	PostalCode	Phone	Fax	Email	SupportRepId
1	Luís	Gonçalves	Embraer - Empresa Brasileira de Aeronáutica S.A.	Av. Brigadeiro Faria Lima, 2170	São José dos Campos	SP	Brazil	12227-000	+55 (12) 3923-5555	+55 (12) 3923-5566	[email protected]	3
2	Leonie	Köhler	None	Theodor-Heuss-Straße 34	Stuttgart	None	Germany	70174	+49 0711 2842222	None	[email protected]	5
3	François	Tremblay	None	1498 rue Bélanger	Montréal	QC	Canada	H2G 1A7	+1 (514) 721-4711	None	[email protected]	3
*/


CREATE TABLE "Employee" (
	"EmployeeId" INTEGER NOT NULL,
	"LastName" NVARCHAR(20) NOT NULL,
	"FirstName" NVARCHAR(20) NOT NULL,
	"Title" NVARCHAR(30),
	"ReportsTo" INTEGER,
	"BirthDate" DATETIME,
	"HireDate" DATETIME,
	"Address" NVARCHAR(70),
	"City" NVARCHAR(40),
	"State" NVARCHAR(40),
	"Country" NVARCHAR(40),
	"PostalCode" NVARCHAR(10),
	"Phone" NVARCHAR(24),
	"Fax" NVARCHAR(24),
	"Email" NVARCHAR(60),
	PRIMARY KEY ("EmployeeId"),
	FOREIGN KEY("ReportsTo") REFERENCES "Employee" ("EmployeeId")
)

/*
Employee 表中的前 3 行数据：
EmployeeId	LastName	FirstName	Title	ReportsTo	BirthDate	HireDate	Address	City	State	Country	PostalCode	Phone	Fax	Email
1	Adams	Andrew	General Manager	None	1962-02-18 00:00:00	2002-08-14 00:00:00	11120 Jasper Ave NW	Edmonton	AB	Canada	T5K 2N1	+1 (780) 428-9482	+1 (780) 428-3457	[email protected]
2	Edwards	Nancy	Sales Manager	1	1958-12-08 00:00:00	2002-05-01 00:00:00	825 8 Ave SW	Calgary	AB	Canada	T2P 2T3	+1 (403) 262-3443	+1 (403) 262-3322	[email protected]
3	Peacock	Jane	Sales Support Agent	2	1973-08-29 00:00:00	2002-04-01 00:00:00	1111 6 Ave SW	Calgary	AB	Canada	T2P 5M5	+1 (403) 262-3443	+1 (403) 262-6712	[email protected]
*/


CREATE TABLE "Genre" (
	"GenreId" INTEGER NOT NULL,
	"Name" NVARCHAR(120),
	PRIMARY KEY ("GenreId")
)

/*
Genre 表中的前 3 行数据：
GenreId	Name
1	Rock
2	Jazz
3	Metal
*/


CREATE TABLE "Invoice" (
	"InvoiceId" INTEGER NOT NULL,
	"CustomerId" INTEGER NOT NULL,
	"InvoiceDate" DATETIME NOT NULL,
	"BillingAddress" NVARCHAR(70),
	"BillingCity" NVARCHAR(40),
	"BillingState" NVARCHAR(40),
	"BillingCountry" NVARCHAR(40),
	"BillingPostalCode" NVARCHAR(10),
	"Total" NUMERIC(10, 2) NOT NULL,
	PRIMARY KEY ("InvoiceId"),
	FOREIGN KEY("CustomerId") REFERENCES "Customer" ("CustomerId")
)

/*
Invoice 表中的前 3 行数据：
InvoiceId	CustomerId	InvoiceDate	BillingAddress	BillingCity	BillingState	BillingCountry	BillingPostalCode	Total
1	2	2009-01-01 00:00:00	Theodor-Heuss-Straße 34	Stuttgart	None	Germany	70174	1.98
2	4	2009-01-02 00:00:00	Ullevålsveien 14	Oslo	None	Norway	0171	3.96
3	8	2009-01-03 00:00:00	Grétrystraat 63	Brussels	None	Belgium	1000	5.94
*/


CREATE TABLE "InvoiceLine" (
	"InvoiceLineId" INTEGER NOT NULL,
	"InvoiceId" INTEGER NOT NULL,
	"TrackId" INTEGER NOT NULL,
	"UnitPrice" NUMERIC(10, 2) NOT NULL,
	"Quantity" INTEGER NOT NULL,
	PRIMARY KEY ("InvoiceLineId"),
	FOREIGN KEY("TrackId") REFERENCES "Track" ("TrackId"),
	FOREIGN KEY("InvoiceId") REFERENCES "Invoice" ("InvoiceId")
)

/*
InvoiceLine 表中的前 3 行数据：
InvoiceLineId	InvoiceId	TrackId	UnitPrice	Quantity
1	1	2	0.99	1
2	1	4	0.99	1
3	2	6	0.99	1
*/


CREATE TABLE "MediaType" (
	"MediaTypeId" INTEGER NOT NULL,
	"Name" NVARCHAR(120),
	PRIMARY KEY ("MediaTypeId")
)

/*
MediaType 表中的前 3 行数据：
MediaTypeId	Name
1	MPEG audio file
2	Protected AAC audio file
3	Protected MPEG-4 video file
*/


CREATE TABLE "Playlist" (
	"PlaylistId" INTEGER NOT NULL,
	"Name" NVARCHAR(120),
	PRIMARY KEY ("PlaylistId")
)

/*
Playlist 表中的前 3 行数据：
PlaylistId	Name
1	Music
2	Movies
3	TV Shows
*/


CREATE TABLE "PlaylistTrack" (
	"PlaylistId" INTEGER NOT NULL,
	"TrackId" INTEGER NOT NULL,
	PRIMARY KEY ("PlaylistId", "TrackId"),
	FOREIGN KEY("TrackId") REFERENCES "Track" ("TrackId"),
	FOREIGN KEY("PlaylistId") REFERENCES "Playlist" ("PlaylistId")
)

/*
PlaylistTrack 表中的前 3 行数据：
PlaylistId	TrackId
1	3402
1	3389
1	3390
*/


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("MediaTypeId") REFERENCES "MediaType" ("MediaTypeId"),
	FOREIGN KEY("GenreId") REFERENCES "Genre" ("GenreId"),
	FOREIGN KEY("AlbumId") REFERENCES "Album" ("AlbumId")
)

/*
Track 表中的前 3 行数据：
TrackId	Name	AlbumId	MediaTypeId	GenreId	Composer	Milliseconds	Bytes	UnitPrice
1	For Those About To Rock (We Salute You)	1	1	1	Angus Young, Malcolm Young, Brian Johnson	343719	11170334	0.99
2	Balls to the Wall	2	2	1	None	342562	5510424	0.99
3	Fast As a Shark	3	2	1	F. Baltes, S. Kaufman, U. Dirkscneider & W. Hoffman	230619	3990994	0.99
*/

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

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

# 用于 studio 的 sql_agent.py
from langchain.agents import create_agent
from langchain_core.messages import SystemMessage


# 初始化一个 LLM
from langchain.chat_models import init_chat_model

llm = init_chat_model("openai:gpt-5")

# 获取数据库，并将其本地存储
import requests, 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}")

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;")}')

SCHEMA = db.get_table_info()

import re
from langchain_core.tools import tool
DENY_RE = re.compile(r"\b(INSERT|UPDATE|DELETE|ALTER|DROP|CREATE|REPLACE|TRUNCATE)\b", re.I)
HAS_LIMIT_TAIL_RE = re.compile(r"(?is)\blimit\b\s+\d+(\s*,\s*\d+)?\s*;?\s*$")

def _safe_sql(q: str) -> str:
    # 标准化
    q = q.strip()
    # 禁止多条语句（允许一个可选的末尾分号）
    if q.count(";") > 1 or (q.endswith(";") and ";" in q[:-1]):
        return "错误：不允许包含多条语句。"
    q = q.rstrip(";").strip()

    # 只读限制
    if not q.lower().startswith("select"):
        return "错误：仅允许 SELECT 语句。"
    if DENY_RE.search(q):
        return "错误：检测到 DML/DDL。仅允许只读查询。"

    # 仅当末尾不存在 LIMIT 时才追加（对空白符/换行符具有鲁棒性）
    if not HAS_LIMIT_TAIL_RE.search(q):
        q += " LIMIT 5"
    return q

@tool
def execute_sql(query: str) -> str:
    """执行一个只读的 SQLite SELECT 查询并返回结果。"""
    query = _safe_sql(query)
    q = query
    if q.startswith("Error:"):
        return q
    try:
        return db.run(q)
    except Exception as e:
        return f"错误：{e}"


SYSTEM = f"""你是一个谨慎的 SQLite 分析师。

权威模式（请勿虚构列或表）：
{SCHEMA}

规则：
- 逐步思考。
- 当你需要数据时，请使用工具 `execute_sql` 调用一个 SELECT 查询。
- 仅限只读；禁止 INSERT/UPDATE/DELETE/ALTER/DROP/CREATE/REPLACE/TRUNCATE。
- 除非用户明确要求，否则结果限制为 5 行。
- 如果工具返回 'Error:'，请修改 SQL 并重试。
- 重试次数最多为 5 次。
- 若 5 次尝试后仍未成功，请向用户返回说明。
- 优先使用明确的列名列表；避免使用 SELECT *。
"""

from langchain.agents import create_agent
from langchain_core.messages import SystemMessage
agent = create_agent(
    model=llm,
    tools=[execute_sql],
    prompt=SystemMessage(content=SYSTEM),
)

# --- 策略配置 ------------------------------------------------------

# 客户允许读取的表
CUSTOMER_ALLOWLIST = {
    "invoice",
    "invoiceline",
    "track",
    "album",
    "artist",
    "genre",
    "mediatype",
    "playlist",
    "playlisttrack",
}

# 需按客户范围限定的表（必须包含 CustomerId = :customer_id）
CUSTOMER_SCOPED = {"invoice", "invoiceline"}

# --- 安全正则表达式 ------------------------------------------------------------
DENY_RE = re.compile(r"\b(INSERT|UPDATE|DELETE|ALTER|DROP|CREATE|REPLACE|TRUNCATE)\b", re.I)
HAS_LIMIT_TAIL_RE = re.compile(r"(?is)\blimit\b\s+\d+(\s*,\s*\d+)?\s*;?\s*$")

# 禁止非纯 SELECT 结构，以简化验证
NON_PLAIN_SQL_RE = re.compile(r"\b(with|union|intersect|except)\b|\(\s*select\b", re.I)

# 提取 FROM/JOIN 中的表名及别名（轻量级解析）
FROM_RE = re.compile(r"\bfrom\s+([\"`\[]?\w+[\"`\]]?)(?:\s+as\s+(\w+)|\s+(\w+))?", re.I)
JOIN_RE = re.compile(r"\bjoin\s+([\"`\[]?\w+[\"`\]]?)(?:\s+as\s+(\w+)|\s+(\w+))?", re.I)

# 对 CustomerId 使用的简单检查
CUSTID_PLACEHOLDER_EQ_RE = re.compile(r"\b(?:\w+\.)?customerid\s*=\s*:customer_id\b", re.I)
CUSTID_NUMERIC_EQ_RE     = re.compile(r"\b(?:\w+\.)?customerid\s*=\s*\d+\b", re.I)


def _normalize_ident(name: str) -> str:
    # 去除引号/反引号/方括号并转为小写
    return re.sub(r'^[\"`\[]|[\"`\]]$', '', name).lower()


def _extract_tables_and_aliases(q: str):
    tables = set()
    alias_map = {}  # 别名 -> 基础表名（小写）
    for m in FROM_RE.finditer(q):
        base = _normalize_ident(m.group(1))
        alias = (m.group(2) or m.group(3) or "").lower()
        tables.add(base)
        if alias:
            alias_map[alias] = base
    for m in JOIN_RE.finditer(q):
        base = _normalize_ident(m.group(1))
        alias = (m.group(2) or m.group(3) or "").lower()
        tables.add(base)
        if alias:
            alias_map[alias] = base
    return tables, alias_map


def _safe_sql(q: str, customer_id: int) -> str:
    # 标准化
    q = q.strip()
    # 禁止多语句（允许末尾一个可选分号）
    if q.count(";") > 1 or (q.endswith(";") and ";" in q[:-1]):
        return "错误：不允许使用多个语句。"
    q = q.rstrip(";").strip()

    # 只读限制
    if not q.lower().startswith("select"):
        return "错误：仅允许 SELECT 语句。"
    if DENY_RE.search(q):
        return "错误：检测到 DML/DDL 语句。仅允许只读查询。"

    # 仅允许纯 SELECT（不允许 CTE、子查询、UNION/INTERSECT/EXCEPT）
    if NON_PLAIN_SQL_RE.search(q):
        return "错误：仅允许纯 SELECT（不允许 CTE/子查询/UNION/INTERSECT/EXCEPT）。"

    # 收集引用的表和别名
    tables, alias_map = _extract_tables_and_aliases(q)
    if not tables:
        return "错误：无法确定引用的表。"

    # 白名单检查
    disallowed = {t for t in tables if t not in CUSTOMER_ALLOWLIST}
    if disallowed:
        bad = ", ".join(sorted(disallowed))
        return f"错误：不允许访问表 [{bad}]。"

    # 客户范围检查
    needs_customer_filter = bool(CUSTOMER_SCOPED & tables)
    if needs_customer_filter:
        # 禁止在 CustomerId 中使用数字字面量
        if CUSTID_NUMERIC_EQ_RE.search(q):
            return "错误：请使用 :customer_id 占位符（禁止使用数字字面量）作为 CustomerId。"

        # 查询文本中必须包含 CustomerId = :customer_id 条件
        if not CUSTID_PLACEHOLDER_EQ_RE.search(q):
            return "错误：涉及 Invoice/InvoiceLine 的查询必须包含 CustomerId = :customer_id 条件。"

        # InvoiceLine 的特殊规则：必须关联 Invoice 表
        if "invoiceline" in tables and "invoice" not in tables:
            return "错误：引用 InvoiceLine 的查询必须关联 Invoice 表并按 CustomerId = :customer_id 过滤。"

    # 若末尾无 LIMIT，则追加 LIMIT 5（兼容空格/换行）
    if not HAS_LIMIT_TAIL_RE.search(q):
        q += " LIMIT 5"
    return q