Build an AI-Powered Data Q&A Bot with DuckDB: Let Business Users Query Data with Natural Language

import duckdb
import json

# Create an in-memory DuckDB database (zero configuration)
con = duckdb.connect(':memory:')

# Directly read CSV with automatic schema inference
con.execute("""
CREATE TABLE orders AS
SELECT * FROM read_csv_auto('https://github.com/duckdb/duckdb-data/raw/main/data/orders.csv');
""")

# Inspect table structure and row count
print(con.execute("DESCRIBE orders").fetchall())
print(con.execute("SELECT COUNT(*) FROM orders").fetchone())

# Also supports Parquet, JSON, Excel, and many other formats
# 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):
    """Function to get database table structure"""
    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": """You are a SQL expert. Users will ask questions in natural language, and you need to convert them into DuckDB SQL queries.
Important rules:
1. Only use the provided tables, do not create new ones
2. Limit query results to 100 rows
3. Use ROUND to keep 2 decimal places for numerical calculations
4. Output only SQL, no explanatory text"""
        },
        {
            "role": "user",
            "content": f"""Database table structure:
{schema_str}

User question: {user_query}

Please generate the SQL query:"""
        }
    ]
    
    response = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=messages,
        temperature=0.1,  # Low temperature for stability
        max_tokens=500
    )
    
    sql = response.choices[0].message.content.strip()
    # Clean markdown code block markers
    if sql.startswith("```"):
        sql = sql.split("\n", 1)[1].rsplit("```", 1)[0].strip()
    return sql

class SafeQueryEngine:
    """Secure DuckDB query executor"""
    
    def __init__(self, connection):
        self.con = connection
    
    def execute_query(self, sql, max_rows=100):
        # Security check: only allow SELECT
        sql_upper = sql.strip().upper()
        if not sql_upper.startswith("SELECT"):
            raise PermissionError("Only SELECT queries are allowed")
        
        dangerous_keywords = ['DROP', 'DELETE', 'INSERT', 'UPDATE', 'ALTER', 'CREATE', 'EXEC']
        for keyword in dangerous_keywords:
            if keyword in sql_upper:
                raise PermissionError(f"Queries containing {keyword} are not allowed")
        
        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"Sorry, the query failed: {query_result['error']}. Please try rephrasing your question."
    
    if query_result['row_count'] == 0:
        return "Based on the current data, no matching results were found."
    
    # Take only the first 5 rows as context to save tokens
    data_summary = json.dumps(query_result['data'][:5], indent=2, ensure_ascii=False)
    
    messages = [
        {
            "role": "system",
            "content": """You are a data analyst. Convert the query results into a Chinese answer that is easy for users to understand.
1. Summarize the core finding (1-2 sentences)
2. List key data points
3. Provide a brief business recommendation"""
        },
        {
            "role": "user",
            "content": f"""User question: {user_query}

Query results (first 5 rows):
{data_summary}

Total records: {query_result['row_count']}

Please answer the user's question in natural language:"""
        }
    ]
    
    response = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=messages,
        temperature=0.3,  # Slightly higher temperature for more natural responses
        max_tokens=1000
    )
    
    return response.choices[0].message.content

def chat_loop():
    print("=" * 60)
    print("🤖 DuckDB Smart Data Q&A Assistant")
    print("=" * 60)
    
    schema_info = get_table_schema(con)
    engine = SafeQueryEngine(con)
    
    while True:
        user_query = input("💬 You: ").strip()
        if user_query.lower() in ('quit', 'exit', 'q'):
            break
        
        # 1. Generate SQL
        sql = generate_sql(user_query, schema_info)
        print(f"📝 Generated SQL: {sql}")
        
        # 2. Execute query
        query_result = engine.execute_query(sql)
        
        # 3. Generate answer
        answer = generate_answer(user_query, query_result, schema_info)
        print(f"\n🤖 Assistant: {answer}")
        print("-" * 60)

# Start the loop
chat_loop()

from fastapi import FastAPI
from pydantic import BaseModel

app = FastAPI(title="DuckDB Data Q&A 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")  # Fully utilize multi-core CPUs

from functools import lru_cache

@lru_cache(maxsize=128)
def cached_ask(question_hash: str, question: str, schema_hash: str):
    """Cache answers for identical questions to avoid repeated LLM calls"""
    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)
            # Try executing to verify validity
            test_result = engine.execute_query(sql)
            if test_result['success']:
                return sql
        except:
            pass
    return None