Update app.py

app.py

@@ -6,13 +6,10 @@ import re
 import os
 from io import BytesIO
 
-# 🔐 Together API Key (set this securely in Hugging Face Spaces)
-TOGETHER_API_KEY = os.getenv("TOGETHER_API_KEY")  # Must be set in the environment variables
+# 🔐 Read Together API key from Hugging Face Secrets
+TOGETHER_API_KEY = os.getenv("TOGETHER_API_KEY")
 
-# 🔧 Global DataFrame holder
-df = None
-
-# 🧠 Generate SQL using Together API
+# 🧠 Generate SQL from Prompt
 def generate_sql_from_prompt(prompt, df):
     schema = ", ".join([f"{col} ({str(dtype)})" for col, dtype in df.dtypes.items()])
     full_prompt = f"""
@@ -36,10 +33,13 @@ Write a valid SQL query using the 'df' table. Return only the SQL code.
         "max_tokens": 200
     }
 
-    response = requests.post(url, headers=headers, json=payload)
-    response.raise_for_status()
-    result = response.json()
-    return result['choices'][0]['message']['content'].strip("```sql").strip("```").strip()
+    try:
+        response = requests.post(url, headers=headers, json=payload)
+        response.raise_for_status()
+        result = response.json()
+        return result['choices'][0]['message']['content'].strip("```sql").strip("```").strip()
+    except Exception as e:
+        return f"Error in API call: {str(e)}"
 
 # 🧽 Clean SQL for DuckDB
 def clean_sql_for_duckdb(sql, df_columns):
@@ -50,42 +50,52 @@ def clean_sql_for_duckdb(sql, df_columns):
             sql = re.sub(pattern, f'"{col}"', sql)
     return sql
 
-# 📦 Handle Excel upload
+# 🔄 Global dataframe holder
+df_global = {"df": None}
+
+# 📂 Upload handler
 def upload_excel(file):
-    global df
-    df = pd.read_excel(BytesIO(file.read()))
-    return f"✅ Loaded file with shape {df.shape}"
+    if file is None:
+        return "❌ No file uploaded."
+    try:
+        df = pd.read_excel(BytesIO(file.read()))
+        df_global["df"] = df
+        return f"✅ Uploaded file with shape: {df.shape}"
+    except Exception as e:
+        return f"❌ Failed to load file: {e}"
 
-# 💬 Handle chat prompt
-def handle_chat(prompt):
-    global df
+# 💬 Handle prompt
+def handle_query(prompt):
+    df = df_global.get("df")
     if df is None:
         return "❌ Please upload an Excel file first.", pd.DataFrame()
 
     try:
         sql = generate_sql_from_prompt(prompt, df)
+        if sql.startswith("Error"):
+            return sql, pd.DataFrame()
+
         cleaned_sql = clean_sql_for_duckdb(sql, df.columns)
         result_df = duckdb.query(cleaned_sql).to_df()
         return f"📜 Generated SQL:\n{sql}", result_df
     except Exception as e:
         return f"❌ Error: {e}", pd.DataFrame()
 
-# 🎨 Gradio UI
+# 🎨 UI
 with gr.Blocks() as demo:
     gr.Markdown("# 🤖 SQL Chatbot with Together API + DuckDB")
-    file_uploader = gr.File(label="📂 Upload Excel File", file_types=[".xlsx"])
-    upload_status = gr.Textbox(label="📄 File Status", interactive=False)
+    file_input = gr.File(label="📂 Upload Excel File (.xlsx only)", file_types=[".xlsx"])
+    upload_status = gr.Textbox(label="Status", interactive=False)
 
     with gr.Row():
-        prompt_box = gr.Textbox(label="💬 Ask your Question", placeholder="e.g., What is the total revenue?")
-        generate_btn = gr.Button("🚀 Generate SQL")
+        prompt_box = gr.Textbox(label="💬 Your Question", placeholder="e.g., Show me total sales by region")
+        run_button = gr.Button("🚀 Generate SQL + Run")
 
     sql_output = gr.Textbox(label="📜 SQL Query")
-    result_table = gr.Dataframe(label="📊 Query Result")
+    result_table = gr.Dataframe(label="📊 Query Results")
 
-    file_uploader.change(upload_excel, inputs=file_uploader, outputs=upload_status)
-    generate_btn.click(fn=handle_chat, inputs=prompt_box, outputs=[sql_output, result_table])
+    file_input.change(upload_excel, inputs=file_input, outputs=upload_status)
+    run_button.click(handle_query, inputs=prompt_box, outputs=[sql_output, result_table])
 
-# 🏁 Launch app
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()