Text-to-SQL 自然语言查询实战

用户输入：上个月 GMV 最高的品类是哪几个？
                    ↓
         [理解意图 + 理解表结构]
                    ↓
SELECT category, SUM(gmv) AS total_gmv
FROM orders
WHERE order_date >= DATE_TRUNC('month', CURRENT_DATE - INTERVAL '1 month')
  AND order_date < DATE_TRUNC('month', CURRENT_DATE)
GROUP BY category
ORDER BY total_gmv DESC
LIMIT 10;

第一步：Schema 检索（过滤不相关的表）
用户问题 → 向量化 → 在表/字段的向量索引中检索 → 筛选 Top-K 相关表

第二步：Schema 注入（格式化后送给 LLM）
相关表的 DDL + 字段注释 + 数据样例 → 结构化 prompt

-- 订单主表，记录所有交易订单
CREATE TABLE orders (
    order_id    BIGINT COMMENT '订单ID，全局唯一',
    user_id     BIGINT COMMENT '用户ID，关联users表',
    category    VARCHAR(50) COMMENT '商品品类，如：服饰、3C、食品',
    gmv         DECIMAL(18,2) COMMENT '订单金额，含税，不含退款',
    order_date  DATE COMMENT '下单日期',
    status      VARCHAR(20) COMMENT '订单状态：paid/refunded/cancelled'
);

[系统角色]
你是一个数据分析助手，帮助将自然语言问题转换为 SQL 查询。

[数据库信息]
数据库类型：Hive SQL（Spark 兼容）
数据库名：dw_retail

[相关表结构]
<这里是筛选后的 DDL + 注释>

[业务规则]（可选，高价值）
- "活跃用户"定义：近 30 天有下单行为的用户
- GMV 统计口径：paid 状态的订单，不含退款
- 分区字段 dt 必须在 WHERE 条件中出现

[历史示例]（Few-shot）
问题：上周各品类的销售额？
SQL：SELECT category, SUM(gmv) ... WHERE dt BETWEEN ...

[当前问题]
用户问题：<用户输入>

[输出要求]
只输出 SQL，不需要解释。SQL 末尾加分号。

# 示意结构
历史 SQL 库 = [
    {
        "question": "上个月各渠道的新增用户数",
        "sql": "SELECT channel, COUNT(DISTINCT user_id) ...",
        "tags": ["用户增长", "渠道分析"]
    },
    ...
]

# 查询时
相关示例 = vector_search(用户问题, 历史SQL库, top_k=3)
prompt = build_prompt(schema, 相关示例, 用户问题)

"""
Text-to-SQL 原型实现
使用 Claude API + DuckDB（可替换为任意数据库）
"""

import anthropic
import duckdb
import json
from typing import Optional

# ============================================================
# 1. 数据库连接和 Schema 提取
# ============================================================

def get_schema_ddl(conn: duckdb.DuckDBPyConnection) -> str:
    """从数据库中提取所有表的 DDL 信息"""
    tables = conn.execute("SHOW TABLES").fetchall()
    ddl_parts = []

    for (table_name,) in tables:
        # 获取建表语句
        create_sql = conn.execute(
            f"SELECT sql FROM sqlite_master WHERE name='{table_name}'"
        ).fetchone()

        # 获取字段信息
        columns = conn.execute(f"DESCRIBE {table_name}").fetchall()
        col_info = "\n".join(
            f"  -- {col[0]}: {col[1]}" for col in columns
        )

        ddl_parts.append(f"-- 表: {table_name}\n{col_info}")

    return "\n\n".join(ddl_parts)


def get_sample_data(conn: duckdb.DuckDBPyConnection, table_name: str, n: int = 3) -> str:
    """获取表的样例数据，帮助 LLM 理解数据格式"""
    try:
        rows = conn.execute(f"SELECT * FROM {table_name} LIMIT {n}").fetchdf()
        return rows.to_string(index=False)
    except Exception:
        return ""


# ============================================================
# 2. Prompt 构建
# ============================================================

SYSTEM_PROMPT = """你是一个数据查询助手，将用户的自然语言问题转换为 SQL 查询。

规则：
1. 只输出 SQL，不要解释，不要 markdown 代码块
2. SQL 末尾加分号
3. 使用 DuckDB SQL 语法
4. 日期处理使用 CURRENT_DATE
5. 字符串比较使用单引号
6. 如果问题不明确，生成最合理的解释对应的 SQL
"""

def build_prompt(
    schema_info: str,
    question: str,
    few_shot_examples: Optional[list] = None,
    business_rules: Optional[str] = None
) -> str:
    """构建发送给 LLM 的 prompt"""

    parts = []

    # Schema 信息
    parts.append(f"## 数据库 Schema\n\n{schema_info}")

    # 业务规则（如果有）
    if business_rules:
        parts.append(f"## 业务规则\n\n{business_rules}")

    # Few-shot 示例（如果有）
    if few_shot_examples:
        examples_text = "\n\n".join(
            f"问题：{ex['question']}\nSQL：{ex['sql']}"
            for ex in few_shot_examples
        )
        parts.append(f"## 参考示例\n\n{examples_text}")

    # 当前问题
    parts.append(f"## 当前问题\n\n{question}")

    return "\n\n---\n\n".join(parts)


# ============================================================
# 3. LLM 调用
# ============================================================

def generate_sql(
    client: anthropic.Anthropic,
    prompt: str,
    model: str = "claude-opus-4-5"
) -> str:
    """调用 Claude API 生成 SQL"""
    message = client.messages.create(
        model=model,
        max_tokens=1024,
        system=SYSTEM_PROMPT,
        messages=[
            {"role": "user", "content": prompt}
        ]
    )
    return message.content[0].text.strip()


# ============================================================
# 4. SQL 执行与结果返回
# ============================================================

def execute_sql(conn: duckdb.DuckDBPyConnection, sql: str) -> dict:
    """执行 SQL 并返回结果"""
    try:
        result_df = conn.execute(sql).fetchdf()
        return {
            "success": True,
            "sql": sql,
            "rows": len(result_df),
            "data": result_df.to_dict(orient="records"),
            "columns": list(result_df.columns)
        }
    except Exception as e:
        return {
            "success": False,
            "sql": sql,
            "error": str(e)
        }


# ============================================================
# 5. 完整的 Text-to-SQL Pipeline
# ============================================================

def text_to_sql_query(
    question: str,
    conn: duckdb.DuckDBPyConnection,
    client: anthropic.Anthropic,
    few_shot_examples: Optional[list] = None,
    business_rules: Optional[str] = None,
    auto_execute: bool = True
) -> dict:
    """
    完整的 Text-to-SQL 流程

    Args:
        question: 自然语言问题
        conn: 数据库连接
        client: Anthropic 客户端
        few_shot_examples: 历史示例列表
        business_rules: 业务规则描述
        auto_execute: 是否自动执行生成的 SQL

    Returns:
        包含 SQL 和查询结果的字典
    """
    # Step 1: 提取 Schema
    schema_info = get_schema_ddl(conn)

    # Step 2: 构建 Prompt
    prompt = build_prompt(
        schema_info=schema_info,
        question=question,
        few_shot_examples=few_shot_examples,
        business_rules=business_rules
    )

    # Step 3: 生成 SQL
    generated_sql = generate_sql(client, prompt)
    print(f"[生成的 SQL]\n{generated_sql}\n")

    # Step 4: 执行（可选）
    if auto_execute:
        result = execute_sql(conn, generated_sql)
        return result
    else:
        return {"sql": generated_sql, "executed": False}


# ============================================================
# 6. 使用示例
# ============================================================

if __name__ == "__main__":
    # 初始化数据库（使用 DuckDB 内存数据库做演示）
    conn = duckdb.connect(":memory:")

    # 创建示例表
    conn.execute("""
        CREATE TABLE orders (
            order_id   INTEGER,
            user_id    INTEGER,
            category   VARCHAR,
            gmv        DECIMAL(10,2),
            order_date DATE,
            status     VARCHAR
        )
    """)

    # 插入测试数据
    conn.execute("""
        INSERT INTO orders VALUES
        (1, 101, '服饰', 299.00, '2026-02-01', 'paid'),
        (2, 102, '3C',   1299.00, '2026-02-05', 'paid'),
        (3, 101, '食品', 88.00,  '2026-02-10', 'refunded'),
        (4, 103, '服饰', 459.00, '2026-02-15', 'paid'),
        (5, 104, '3C',   3999.00, '2026-02-20', 'paid')
    """)

    # 初始化 Anthropic 客户端
    client = anthropic.Anthropic()  # 从环境变量读取 ANTHROPIC_API_KEY

    # 定义业务规则
    business_rules = """
    - GMV 统计只包含 status='paid' 的订单，不含 refunded
    - 品类字段 category 的有效值：服饰、3C、食品、家居、美妆
    """

    # 定义 Few-shot 示例
    examples = [
        {
            "question": "各品类的订单数量",
            "sql": "SELECT category, COUNT(*) AS order_count FROM orders WHERE status='paid' GROUP BY category ORDER BY order_count DESC;"
        }
    ]

    # 执行查询
    result = text_to_sql_query(
        question="上个月 GMV 最高的品类是哪些？",
        conn=conn,
        client=client,
        few_shot_examples=examples,
        business_rules=business_rules
    )

    if result["success"]:
        print(f"查询成功，返回 {result['rows']} 行数据")
        print(json.dumps(result["data"], ensure_ascii=False, indent=2))
    else:
        print(f"执行失败：{result['error']}")

你：@schema/orders.sql @schema/users.sql

帮我写一个查询：找出近 90 天内，购买过 3C 品类但从未购买过服饰品类的用户，
按他们的 3C 消费总额降序排列，取前 100 名。

Claude：[生成 SQL]

你：把时间条件改成近 30 天，另外加上用户的注册时间

Claude：[修改后的 SQL]

# business_glossary.yaml 示例结构

metrics:
  gmv:
    definition: 成交总额
    sql_expression: "SUM(CASE WHEN status='paid' THEN amount ELSE 0 END)"
    table: orders
    note: 不含退款，不含运费

  dau:
    definition: 日活跃用户数
    sql_expression: "COUNT(DISTINCT user_id)"
    table: user_events
    filter: "event_date = CURRENT_DATE"
    note: 有任意行为即算活跃，包括浏览

dimensions:
  channel:
    field: acquisition_channel
    table: users
    values: [organic, paid_search, social, referral]

  category:
    field: first_category
    table: items
    note: 使用一级品类，不是二级品类