Spaces:
Running
Running
import gradio as gr | |
import openai | |
import sqlite3 | |
import pandas as pd | |
import matplotlib.pyplot as plt | |
import seaborn as sns | |
import os | |
# OpenRouter API Key (Replace with yours) | |
OPENROUTER_API_KEY = "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" | |
OPENROUTER_MODEL = "sophosympatheia/rogue-rose-103b-v0.2:free" | |
# Database Path | |
db_path = "ecommerce.db" | |
# Ensure dataset exists | |
if not os.path.exists(db_path): | |
print("Database file not found! Please upload ecommerce.db.") | |
# Initialize OpenAI client | |
openai_client = openai.OpenAI(api_key=OPENROUTER_API_KEY, base_url="https://openrouter.ai/api/v1") | |
# Updated Few-Shot Examples with SQLite-Compatible Queries | |
few_shot_examples = [ | |
{"input": "Find the busiest months for orders.", | |
"output": "SELECT strftime('%m', order_purchase_timestamp) AS month, COUNT(*) AS order_count FROM orders GROUP BY month ORDER BY order_count DESC;"}, | |
{"input": "Show all customers from São Paulo.", | |
"output": "SELECT * FROM customers WHERE customer_state = 'SP';"}, | |
{"input": "Find the total sales per product.", | |
"output": "SELECT product_id, SUM(price) FROM order_items GROUP BY product_id;"}, | |
{"input": "List all orders placed in 2017.", | |
"output": "SELECT * FROM orders WHERE order_purchase_timestamp LIKE '2017%';"} | |
] | |
# Function: Convert Text to SQL | |
def text_to_sql(query): | |
prompt = "Convert the following queries into SQLite-compatible SQL:\n\n" | |
for example in few_shot_examples: | |
prompt += f"Input: {example['input']}\nOutput: {example['output']}\n\n" | |
prompt += f"Input: {query}\nOutput:" | |
try: | |
response = openai_client.chat.completions.create( | |
model=OPENROUTER_MODEL, | |
messages=[{"role": "system", "content": "You are an SQLite expert."}, | |
{"role": "user", "content": prompt}] | |
) | |
sql_query = response.choices[0].message.content.strip() | |
return sql_query if sql_query.lower().startswith("select") else f"Error: Invalid SQL generated - {sql_query}" | |
except Exception as e: | |
return f"Error: {e}" | |
# Function: Execute SQL on SQLite Database | |
def execute_sql(sql_query): | |
try: | |
conn = sqlite3.connect(db_path) | |
df = pd.read_sql_query(sql_query, conn) | |
conn.close() | |
return df | |
except Exception as e: | |
return f"SQL Execution Error: {e}" | |
# Function: Generate Data Visualization | |
def visualize_data(df): | |
if df.empty or df.shape[1] < 2: | |
return None | |
numeric_cols = df.select_dtypes(include=['number']).columns | |
if len(numeric_cols) < 1: | |
return None | |
plt.figure(figsize=(6, 4)) | |
sns.set_theme(style="darkgrid") | |
if len(numeric_cols) == 1: | |
sns.histplot(df[numeric_cols[0]], bins=10, kde=True, color="teal") | |
plt.title(f"Distribution of {numeric_cols[0]}") | |
elif len(numeric_cols) == 2: | |
sns.scatterplot(x=df[numeric_cols[0]], y=df[numeric_cols[1]], color="blue") | |
plt.title(f"{numeric_cols[0]} vs {numeric_cols[1]}") | |
elif df.shape[0] < 10: | |
plt.pie(df[numeric_cols[0]], labels=df.iloc[:, 0], autopct='%1.1f%%', colors=sns.color_palette("pastel")) | |
plt.title(f"Proportion of {numeric_cols[0]}") | |
else: | |
sns.barplot(x=df.iloc[:, 0], y=df[numeric_cols[0]], palette="coolwarm") | |
plt.xticks(rotation=45) | |
plt.title(f"{df.columns[0]} vs {numeric_cols[0]}") | |
plt.tight_layout() | |
plt.savefig("chart.png") | |
return "chart.png" | |
# Gradio UI | |
def gradio_ui(query): | |
sql_query = text_to_sql(query) | |
results = execute_sql(sql_query) | |
visualization = visualize_data(results) if isinstance(results, pd.DataFrame) else None | |
return sql_query, results.to_string(index=False) if isinstance(results, pd.DataFrame) else results, visualization | |
with gr.Blocks() as demo: | |
gr.Markdown("## SQL Explorer: Text to SQL with a Simple Visualization") | |
query_input = gr.Textbox(label="Enter your query", placeholder="Enter your query in English.") | |
submit_btn = gr.Button("Convert & Execute") | |
sql_output = gr.Textbox(label="Generated SQL Query") | |
table_output = gr.Textbox(label="Query Results") | |
chart_output = gr.Image(label="Data Visualization") | |
submit_btn.click(gradio_ui, inputs=[query_input], outputs=[sql_output, table_output, chart_output]) | |
# Launch | |
demo.launch() | |