文本到 SQL 生成：自然语言查询数据库

bash pip install openai sqlite3

本示例使用 OpenAI API 调用大语言模型，数据库采用 SQLite。

### 第一步：设计样例数据库

创建一个包含三张表的简单电商数据库。

```sql
-- 文件：schema.sql
CREATE TABLE customers (
    customer_id INTEGER PRIMARY KEY,
    name TEXT,
    city TEXT
);

CREATE TABLE products (
    product_id INTEGER PRIMARY KEY,
    product_name TEXT,
    category TEXT,
    price REAL
);

CREATE TABLE orders (
    order_id INTEGER PRIMARY KEY,
    customer_id INTEGER,
    product_id INTEGER,
    quantity INTEGER,
    order_date DATE,
    FOREIGN KEY(customer_id) REFERENCES customers(customer_id),
    FOREIGN KEY(product_id) REFERENCES products(product_id)
);

插入少量测试数据：

INSERT INTO customers VALUES (1,'Alice','New York'),(2,'Bob','Los Angeles');
INSERT INTO products VALUES (1,'Laptop','Electronics',1200),(2,'Headphones','Electronics',150);
INSERT INTO orders VALUES (1,1,1,1,'2025-03-10'),(2,2,2,2,'2025-03-12');

第二步：构造大模型提示词

将数据库 schema 和用户问题组装，通过 few-shot 示例引导模型。

# text2sql.py
import openai
import sqlite3

openai.api_key = "your-api-key"

schema = """
表: customers
列: customer_id, name, city

表: products
列: product_id, product_name, category, price

表: orders
列: order_id, customer_id, product_id, quantity, order_date
外键: customer_id 引用 customers, product_id 引用 products
"""

def build_prompt(question):
    prompt = f"""你是一个SQL查询生成助手。根据以下数据库结构，将用户问题转换为准确的SQL语句。
{schema}

示例:
问题: 列出所有客户姓名。
SQL: SELECT name FROM customers;

问题: 2025年3月10日之后的所有订单总数是多少？
SQL: SELECT COUNT(*) FROM orders WHERE order_date > '2025-03-10';

现在请转换:
问题: {question}
SQL: """
    return prompt

def generate_sql(question):
    response = openai.ChatCompletion.create(
        model="gpt-3.5-turbo",
        messages=[{"role": "user", "content": build_prompt(question)}],
        temperature=0
    )
    return response.choices[0].message.content.strip()

第三步：执行并呈现结果

def run_query(sql):
    conn = sqlite3.connect("ecommerce.db")
    cur = conn.cursor()
    cur.execute(sql)
    rows = cur.fetchall()
    conn.close()
    return rows

# 用户提问
user_question = "哪个客户买了最贵的产品？"
sql = generate_sql(user_question)
print(f"生成的SQL: {sql}")

# 安全起见，可增加简单的 SQL 执行检查，例如只允许 SELECT
if sql.strip().lower().startswith("select"):
    results = run_query(sql)
    print("查询结果:", results)
else:
    print("非 SELECT 语句，已阻止执行")

可能的输出

生成的SQL: SELECT customers.name FROM customers 
JOIN orders ON customers.customer_id = orders.customer_id 
JOIN products ON orders.product_id = products.product_id 
ORDER BY products.price DESC LIMIT 1;
查询结果: [('Alice',)]