Update app.py

app.py

@@ -10,27 +10,26 @@ import os
 OPENROUTER_API_KEY = "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"
 OPENROUTER_MODEL = "sophosympatheia/rogue-rose-103b-v0.2:free"
 
-# Path to the SQLite Database
+# Hugging Face Space path
 DB_PATH = "ecommerce.db"
 
-# Ensure database exists before proceeding
+# Ensure dataset exists
 if not os.path.exists(DB_PATH):
-    raise FileNotFoundError("Database file 'ecommerce.db' not found. Upload it first!")
+    os.system("wget https://your-dataset-link.com/ecommerce.db -O ecommerce.db")  # Replace with actual dataset link
 
 # Initialize OpenAI client
 openai_client = openai.OpenAI(api_key=OPENROUTER_API_KEY, base_url="https://openrouter.ai/api/v1")
 
-# Few-shot examples with SQLite-friendly queries
+# 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) AS total_sales FROM order_items GROUP BY product_id;"},
-    {"input": "List all orders placed in 2017.", "output": "SELECT * FROM orders WHERE order_purchase_timestamp LIKE '2017%';"},
-    {"input": "Find the busiest months for orders.", "output": "SELECT strftime('%m', order_purchase_timestamp) AS order_month, COUNT(*) AS orders_count FROM orders GROUP BY order_month ORDER BY orders_count DESC;"},
+    {"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"
+    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:"
@@ -40,13 +39,18 @@ def text_to_sql(query):
             model=OPENROUTER_MODEL,
             messages=[{"role": "system", "content": "You are an SQL expert."}, {"role": "user", "content": prompt}]
         )
-        return response.choices[0].message.content.strip()
+        sql_query = response.choices[0].message.content.strip()
+        sql_query = sql_query.split("\n")[0]  # Take only the first line if multiple lines exist
+        sql_query = sql_query.replace("mathchar", "").rstrip(";")  # Remove unwanted text
+        return sql_query
     except Exception as e:
         return f"Error: {e}"
 
+# Function: Execute SQL on SQLite database
 def execute_sql(sql_query):
     try:
         sql_query = sql_query.strip().rstrip(";")  # Remove trailing semicolons
+        sql_query = sql_query.replace("mathchar", "")  # Remove any bad tokens
         conn = sqlite3.connect(DB_PATH)
         df = pd.read_sql_query(sql_query, conn)
         conn.close()
@@ -54,7 +58,6 @@ def execute_sql(sql_query):
     except Exception as e:
         return f"SQL Execution Error: {e}"
 
-
 # Function: Generate Dynamic Visualization
 def visualize_data(df):
     if df.empty or df.shape[1] < 2:
@@ -96,8 +99,8 @@ def gradio_ui(query):
     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 Real Execution & Visualization")
-    query_input = gr.Textbox(label="Enter your query", placeholder="e.g., Show all products sold in 2018.")
+    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")