import duckdb
import json

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

# Simulate e-commerce order data
con.execute("""
CREATE TABLE orders AS
SELECT * FROM read_csv_auto('https://github.com/duckdb/duckdb-data/raw/main/data/orders.csv');
""")

# View table structure
print(con.execute("DESCRIBE orders").fetchdf())

# Read local CSV (auto-infer schema)
con.read_csv_auto('/path/to/sales_data.csv')

# Read Parquet files (columnar storage, faster)
con.read_parquet('/path/to/sales_data.parquet')

# Batch-read all Parquet files in a directory
con.read_parquet('/path/to/data/*.parquet')

import os
from openai import OpenAI

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

def get_table_schema(con):
 """Get schema information for all tables in DuckDB"""
 tables = con.execute("SHOW TABLES").fetchall()
 schema_info = {}
 for table_name, in tables:
 columns = con.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):
 """Translate natural language to SQL"""
 schema_str = json.dumps(schema_info, indent=2, ensure_ascii=False)
 
 messages = [
 {
 "role": "system",
 "content": """You are a SQL expert. The user will ask questions in natural language. 
 You need to convert them into DuckDB SQL queries.
 
 Important rules:
 1. Only use the tables listed below — do not create new tables
 2. Limit results to 100 rows
 3. For numerical calculations, use ROUND to 2 decimal places
 4. If the question cannot be answered with existing data, return an empty string
 5. Output only the SQL, no explanatory text"""
 },
 {
 "role": "user",
 "content": f"""Database table schema:
{schema_str}

User question: {user_query}

Please generate the SQL query (output only the SQL statement):"""
 }
 ]
 
 response = client.chat.completions.create(
 model="gpt-4o-mini",
 messages=messages,
 temperature=0.1,
 max_tokens=500
 )
 
 sql = response.choices[0].message.content.strip()
 # Clean possible 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):
 """Execute SQL and return results"""
 # Security check: only allow SELECT
 sql_upper = sql.strip().upper()
 if not sql_upper.startswith("SELECT"):
 raise PermissionError("Only SELECT queries are allowed")
 
 # Check for dangerous functions
 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()
 
 # Limit rows
 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):
 """Convert query results into natural language answers"""
 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."
 
 # Build a brief data summary
 data_summary = json.dumps(query_result['data'][:5], indent=2, ensure_ascii=False)
 
 messages = [
 {
 "role": "system",
 "content": """You are a data analyst. The user asked a question in natural language. 
 You have executed the SQL query and received results.
 Convert the results into an easy-to-understand answer. Requirements:
 1. Summarize the core finding (1-2 sentences)
 2. List key data points
 3. Provide brief business recommendations
 4. If the dataset is large, show only the first few rows"""
 },
 {
 "role": "user",
 "content": f"""User question: {user_query}

Query results (first 5 rows):
{data_summary}

Total: {query_result['row_count']} records

Please answer the user's question in natural language:"""
 }
 ]
 
 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():
 """Interactive conversation loop"""
 print("=" * 60)
 print("🤖 DuckDB Intelligent Data Q&A Assistant")
 print("=" * 60)
 print("Type your question (type 'quit' to exit)")
 print()
 
 schema_info = get_table_schema(con)
 engine = SafeQueryEngine(con)
 
 while True:
 user_query = input("💬 You: ").strip()
 if user_query.lower() in ('quit', 'exit', 'q'):
 print("Goodbye! 👋")
 break
 
 # Step 1: Generate SQL
 print("\n⏳ Generating query...")
 sql = generate_sql(user_query, schema_info)
 print(f" SQL: {sql}")
 
 if not sql:
 print("❌ Unable to understand your question. Please rephrase.")
 continue
 
 # Step 2: Execute query
 print("⏳ Executing query...")
 query_result = engine.execute_query(sql)
 
 if not query_result['success']:
 print(f"❌ {query_result['error']}")
 continue
 
 # Step 3: Generate natural language answer
 print("⏳ Generating answer...")
 answer = generate_answer(user_query, query_result, schema_info)
 print(f"\n🤖 Assistant: {answer}")
 print()
 print("-" * 60)

# Start the chat
if __name__ == "__main__":
 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}

from functools import lru_cache

@lru_cache(maxsize=1000)
def cached_query(sql_hash, sql_result_json):
 pass

# Enable multi-threading for acceleration
con.execute("SET threads TO 4")

def complex_query_pipeline(user_query, schema_info):
 """Multi-step query decomposition"""
 # Step 1: Let LLM decide which tables and fields are needed
 plan = generate_query_plan(user_query, schema_info)
 # Step 2: Execute each sub-query step by step
 # Step 3: Combine results into final answer
 ...