Update app.py

app.py

@@ -5,6 +5,7 @@ import pandas as pd
 import matplotlib.pyplot as plt
 import seaborn as sns
 import os
+from typing import Optional, Tuple
 
 # OpenRouter API Key (Replace with yours)
 OPENROUTER_API_KEY = "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"
@@ -20,56 +21,65 @@ if not os.path.exists(DB_PATH):
 # Initialize OpenAI client
 openai_client = openai.OpenAI(api_key=OPENROUTER_API_KEY, base_url="https://openrouter.ai/api/v1")
 
-# Few-shot examples for text-to-SQL
-few_shot_examples = [
-    {"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: Fetch database schema
+def fetch_schema(db_path: str) -> str:
+    conn = sqlite3.connect(db_path)
+    cursor = conn.cursor()
+    cursor.execute("SELECT name FROM sqlite_master WHERE type='table';")
+    tables = cursor.fetchall()
+    schema = ""
+    for table in tables:
+        table_name = table[0]
+        cursor.execute(f"PRAGMA table_info({table_name});")
+        columns = cursor.fetchall()
+        schema += f"Table: {table_name}\n"
+        for column in columns:
+            schema += f"  Column: {column[1]}, Type: {column[2]}\n"
+    conn.close()
+    return schema
 
 # Function: Convert text to SQL
-def text_to_sql(query):
-    prompt = "Convert the following queries into SQL:\n\n"
-    for example in few_shot_examples:
-        prompt += f"Input: {example['input']}\nOutput: {example['output']}\n\n"
-    prompt += f"Input: {query}\nOutput:"
-
+def text_to_sql(query: str, schema: str) -> str:
+    prompt = (
+        "You are an SQL expert. Given the following database schema:\n\n"
+        f"{schema}\n\n"
+        "Convert the following query into SQL:\n\n"
+        f"Query: {query}\n"
+        "SQL:"
+    )
     try:
         response = openai_client.chat.completions.create(
             model=OPENROUTER_MODEL,
             messages=[{"role": "system", "content": "You are an SQL expert."}, {"role": "user", "content": prompt}]
         )
         sql_query = response.choices[0].message.content.strip()
-        
-        # Ensure only one query is returned (remove extra text)
-        sql_query = sql_query.split("\n")[0].strip()
         return sql_query
     except Exception as e:
         return f"Error: {e}"
 
 # Function: Execute SQL on SQLite database
-def execute_sql(sql_query):
+def execute_sql(sql_query: str) -> Tuple[Optional[pd.DataFrame], Optional[str]]:
     try:
         conn = sqlite3.connect(DB_PATH)
         df = pd.read_sql_query(sql_query, conn)
         conn.close()
-        return df
+        return df, None
     except Exception as e:
-        return f"SQL Execution Error: {e}"
+        return None, f"SQL Execution Error: {e}"
 
 # Function: Generate Dynamic Visualization
-def visualize_data(df):
+def visualize_data(df: pd.DataFrame) -> Optional[str]:
     if df.empty or df.shape[1] < 2:
         return None
 
+    plt.figure(figsize=(6, 4))
+    sns.set_theme(style="darkgrid")
+
     # Detect numeric columns
     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")
-
     # Choose visualization type dynamically
     if len(numeric_cols) == 1:  # Single numeric column, assume it's a count metric
         sns.histplot(df[numeric_cols[0]], bins=10, kde=True, color="teal")
@@ -90,13 +100,16 @@ def visualize_data(df):
     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
-
+def gradio_ui(query: str) -> Tuple[str, str, Optional[str]]:
+    schema = fetch_schema(DB_PATH)
+    sql_query = text_to_sql(query, schema)
+    df, error = execute_sql(sql_query)
+    if error:
+        return sql_query, error, None
+    visualization = visualize_data(df) if df is not None else None
+    return sql_query, df.to_string(index=False), visualization
+
+# Launch Gradio App
 with gr.Blocks() as demo:
     gr.Markdown("## SQL Explorer: Text-to-SQL with Real Execution & Visualization")
     query_input = gr.Textbox(label="Enter your query", placeholder="e.g., Show all products sold in 2018.")
@@ -107,5 +120,4 @@ with gr.Blocks() as demo:
 
     submit_btn.click(gradio_ui, inputs=[query_input], outputs=[sql_output, table_output, chart_output])
 
-# Launch
-demo.launch()
+demo.launch()