用 Ollama 搭一个 AI SQL 翻译器：自然语言查数据库，不用学 SQL

# 安装 Ollama（如果还没装）
curl -fsSL https://ollama.ai/install.sh | sh

# 拉取模型
ollama pull qwen2.5:7b

# 启动 Ollama 服务
ollama serve &

# 安装 Python 依赖
pip install ollama sqlite3-utils

# setup_db.py
import sqlite3
from datetime import datetime, timedelta

conn = sqlite3.connect("shop.db")
cursor = conn.cursor()

# 创建表
cursor.execute("""
CREATE TABLE IF NOT EXISTS products (
    id INTEGER PRIMARY KEY,
    name TEXT NOT NULL,
    category TEXT,
    price REAL
)
""")

cursor.execute("""
CREATE TABLE IF NOT EXISTS orders (
    id INTEGER PRIMARY KEY,
    product_id INTEGER,
    quantity INTEGER,
    order_date TEXT,
    FOREIGN KEY (product_id) REFERENCES products(id)
)
""")

# 插入示例数据
products = [
    (1, "MacBook Pro", "电脑", 12999),
    (2, "iPhone 15", "手机", 7999),
    (3, "AirPods Pro", "耳机", 1899),
    (4, "iPad Air", "平板", 4799),
    (5, "Apple Watch", "手表", 2999),
    (6, "Magic Keyboard", "配件", 999),
    (7, "Studio Display", "显示器", 11499),
    (8, "Mac Mini", "电脑", 4499),
    (9, "HomePod", "音箱", 1999),
    (10, "AirTag", "配件", 229),
]

orders = []
base_date = datetime(2026, 7, 1)
for i in range(200):
    day_offset = i % 31
    hour = i % 24
    order_date = (base_date + timedelta(days=day_offset)).strftime("%Y-%m-%d")
    product_id = (i % 10) + 1
    quantity = (i % 5) + 1
    orders.append((product_id, quantity, order_date))

cursor.executemany("INSERT INTO products VALUES (?, ?, ?, ?)", products)
cursor.executemany("INSERT INTO orders (product_id, quantity, order_date) VALUES (?, ?, ?)", orders)

conn.commit()
conn.close()
print("✅ 数据库初始化完成：shop.db")

python setup_db.py

# sql_translator.py
import sqlite3
import ollama
import json

MODEL = "qwen2.5:7b"

SYSTEM_PROMPT = """你是一个 SQL 专家。用户用自然语言描述需求，你生成正确的 SQL 查询语句。

规则：
1. 只输出 SQL，不要解释
2. SQL 以分号结尾
3. 如果需求不明确，询问具体字段或条件
4. 只查询不修改数据（SELECT，禁止 INSERT/UPDATE/DELETE）
5. 表结构如下：
   - products: id, name, category, price
   - orders: id, product_id, quantity, order_date
"""

def translate_to_sql(natural_query: str) -> str:
    """自然语言 → SQL"""
    response = ollama.chat(
        model=MODEL,
        messages=[
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": natural_query}
        ]
    )
    sql = response["message"]["content"].strip()
    # 清理可能的 markdown 代码块
    if sql.startswith("```"):
        sql = sql.split("\n", 1)[1].rsplit("```", 1)[0].strip()
    return sql


def execute_sql(db_path: str, sql: str) -> list:
    """执行 SQL 并返回结果"""
    conn = sqlite3.connect(db_path)
    conn.row_factory = sqlite3.Row
    cursor = conn.cursor()
    
    try:
        cursor.execute(sql)
        rows = cursor.fetchall()
        if cursor.description:
            columns = [desc[0] for desc in cursor.description]
            result = [dict(row) for row in rows]
            return {"columns": columns, "data": result, "count": len(result)}
        return {"columns": [], "data": [], "count": 0}
    except Exception as e:
        return {"error": str(e), "sql": sql}
    finally:
        conn.close()


def query(natural_query: str, db_path: str = "shop.db") -> dict:
    """完整流程：自然语言 → SQL → 结果"""
    print(f"\n📝 你的问题: {natural_query}")
    
    # 1. 翻译
    sql = translate_to_sql(natural_query)
    print(f"🤖 生成 SQL: {sql}")
    
    # 2. 执行
    result = execute_sql(db_path, sql)
    
    if "error" in result:
        print(f"❌ 执行失败: {result['error']}")
        return result
    
    print(f"📊 返回 {result['count']} 条记录")
    
    # 3. 格式化输出
    if result["data"]:
        print("\n" + "=" * 50)
        print(" | ".join(result["columns"]))
        print("=" * 50)
        for row in result["data"][:10]:  # 最多显示10条
            print(" | ".join(str(v) for v in row.values()))
        if result["count"] > 10:
            print(f"... 还有 {result['count'] - 10} 条记录")
    
    return result

# main.py
import sys
from sql_translator import query

EXAMPLES = [
    "上个月销售额最高的10个产品",
    "每个类别的产品数量和平均价格",
    "7月份订单数量最多的日期",
    "价格超过5000的产品有哪些",
    "统计每个产品的总销售额",
]

def main():
    if len(sys.argv) > 1:
        # 命令行模式
        natural_query = " ".join(sys.argv[1:])
        result = query(natural_query)
    else:
        # 交互模式
        print("🧠 AI SQL 翻译器")
        print("输入自然语言问题，按 Ctrl+C 退出")
        print("示例问题：")
        for i, ex in enumerate(EXAMPLES, 1):
            print(f"  {i}. {ex}")
        print()
        
        while True:
            try:
                user_input = input("❓ 你的问题: ").strip()
                if not user_input:
                    continue
                if user_input.lower() in ("quit", "exit", "q"):
                    break
                query(user_input)
                print()
            except KeyboardInterrupt:
                print("\n👋 再见！")
                break

if __name__ == "__main__":
    main()

# 先初始化数据库
python setup_db.py

# 交互模式运行
python main.py

# 或直接命令行
python main.py "上个月销售额最高的10个产品"

# 快速开始
ollama pull qwen2.5:7b
pip install ollama
python main.py "你好，查一下今天的订单数"