用 DuckDB 搭建 AI 数据问答机器人：让业务人员自己查数据

import duckdb
import json

# 创建 DuckDB 内存数据库（零配置）
con = duckdb.connect(':memory:')

# 直接读取 CSV（自动推断 schema）
con.execute("""
CREATE TABLE orders AS
SELECT * FROM read_csv_auto('https://github.com/duckdb/duckdb-data/raw/main/data/orders.csv');
""")

# 查看表结构和数据量
print(con.execute("DESCRIBE orders").fetchall())
print(con.execute("SELECT COUNT(*) FROM orders").fetchone())

# 也支持 Parquet、JSON、Excel 等多种格式
# con.read_parquet('/path/to/data/*.parquet')
# con.execute("CREATE TABLE events AS SELECT * FROM read_json_auto('data.json')")

from openai import OpenAI
import os

client = OpenAI(api_key=os.environ.get('OPENAI_API_KEY'))

def get_table_schema(connection):
    """获取数据库表结构的函数"""
    tables = connection.execute(
        "SELECT table_name FROM information_schema.tables WHERE table_schema = 'main'"
    ).fetchall()
    
    schema_info = {}
    for (table_name,) in tables:
        columns = connection.execute(
            f"DESCRIBE {table_name}"
        ).fetchall()
        schema_info[table_name] = [
            {"column": col[0], "type": col[1]} for col in columns
        ]
    return schema_info

def generate_sql(user_query, schema_info):
    schema_str = json.dumps(schema_info, indent=2, ensure_ascii=False)
    
    messages = [
        {
            "role": "system",
            "content": """你是一个 SQL 专家。用户会用自然语言提问，你需要将其转换为 DuckDB SQL 查询。
重要规则：
1. 只允许使用提供的表，不要创建新表
2. 查询结果控制在 100 行以内
3. 如果涉及数值计算，使用 ROUND 保留 2 位小数
4. 只输出 SQL，不要输出解释文字"""
        },
        {
            "role": "user",
            "content": f"""数据库表结构：
{schema_str}

用户问题：{user_query}

请生成 SQL 查询："""
        }
    ]
    
    response = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=messages,
        temperature=0.1,  # 低温度保证稳定性
        max_tokens=500
    )
    
    sql = response.choices[0].message.content.strip()
    # 清理 markdown 代码块标记
    if sql.startswith("```"):
        sql = sql.split("\n", 1)[1].rsplit("```", 1)[0].strip()
    return sql

class SafeQueryEngine:
    """安全的 DuckDB 查询执行器"""
    
    def __init__(self, connection):
        self.con = connection
    
    def execute_query(self, sql, max_rows=100):
        # 安全检查：只允许 SELECT
        sql_upper = sql.strip().upper()
        if not sql_upper.startswith("SELECT"):
            raise PermissionError("只允许执行 SELECT 查询")
        
        dangerous_keywords = ['DROP', 'DELETE', 'INSERT', 'UPDATE', 'ALTER', 'CREATE', 'EXEC']
        for keyword in dangerous_keywords:
            if keyword in sql_upper:
                raise PermissionError(f"不允许执行包含 {keyword} 的查询")
        
        try:
            result = self.con.execute(sql)
            df = result.fetchdf()
            if len(df) > max_rows:
                df = df.head(max_rows)
            return {
                'success': True,
                'data': df.to_dict(orient='records'),
                'columns': list(df.columns),
                'row_count': len(df)
            }
        except Exception as e:
            return {'success': False, 'error': str(e)}

def generate_answer(user_query, query_result, schema_info):
    if not query_result['success']:
        return f"抱歉，查询执行失败：{query_result['error']}。请换一种方式提问。"
    
    if query_result['row_count'] == 0:
        return "根据当前数据，没有找到匹配的结果。"
    
    # 只取前 5 条数据作为上下文，节省 token
    data_summary = json.dumps(query_result['data'][:5], indent=2, ensure_ascii=False)
    
    messages = [
        {
            "role": "system",
            "content": """你是一个数据分析师。请将查询结果转化为用户容易理解的中文回答。
1. 先总结核心发现（1-2句话）
2. 列出关键数据点
3. 给出简短的业务建议"""
        },
        {
            "role": "user",
            "content": f"""用户问题：{user_query}

查询结果（前5条）：
{data_summary}

共 {query_result['row_count']} 条记录

请用自然语言回答用户的问题："""
        }
    ]
    
    response = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=messages,
        temperature=0.3,  # 稍高温度让回答更自然
        max_tokens=1000
    )
    
    return response.choices[0].message.content

def chat_loop():
    print("=" * 60)
    print("🤖 DuckDB 智能数据问答助手")
    print("=" * 60)
    
    schema_info = get_table_schema(con)
    engine = SafeQueryEngine(con)
    
    while True:
        user_query = input("💬 你: ").strip()
        if user_query.lower() in ('quit', 'exit', 'q'):
            break
        
        # 1. 生成 SQL
        sql = generate_sql(user_query, schema_info)
        print(f"📝 生成 SQL: {sql}")
        
        # 2. 执行查询
        query_result = engine.execute_query(sql)
        
        # 3. 生成回答
        answer = generate_answer(user_query, query_result, schema_info)
        print(f"\n🤖 助手: {answer}")
        print("-" * 60)

# 启动
chat_loop()

from fastapi import FastAPI
from pydantic import BaseModel

app = FastAPI(title="DuckDB 数据问答 API")

class QueryRequest(BaseModel):
    question: str

@app.post("/ask")
def ask_question(req: QueryRequest):
    sql = generate_sql(req.question, schema_info)
    result = engine.execute_query(sql)
    answer = generate_answer(req.question, result, schema_info)
    return {"answer": answer}

con.execute("SET threads TO 4")  # 充分利用多核 CPU

from functools import lru_cache

@lru_cache(maxsize=128)
def cached_ask(question_hash: str, question: str, schema_hash: str):
    """缓存相同问题的回答，避免重复调用 LLM"""
    sql = generate_sql(question, schema_info)
    result = engine.execute_query(sql)
    return generate_answer(question, result, schema_info)

def robust_generate_sql(user_query, schema_info, max_retries=3):
    for attempt in range(max_retries):
        try:
            sql = generate_sql(user_query, schema_info)
            # 尝试执行看看是否有效
            test_result = engine.execute_query(sql)
            if test_result['success']:
                return sql
        except:
            pass
    return None