Update streamlit_app.py

streamlit_app.py

@@ -1,54 +1,59 @@
 import streamlit as st
 import pandas as pd
 import requests
-import os
+from pandasql import sqldf
 
-# Page setup
-st.set_page_config(page_title="🧠 Excel SQL Assistant (Together AI)", layout="centered")
-st.title("📊 Excel to SQL with Together AI + Streamlit")
+# 🛠 Page setup
+st.set_page_config(page_title="🧠 Excel SQL Assistant", layout="centered")
+st.title("📊 Excel to SQL with Together AI")
 
-# API key from Hugging Face Secrets
+# 🔐 Load API Key securely
 TOGETHER_API_KEY = st.secrets.get("TOGETHER_API_KEY", None)
 
 if not TOGETHER_API_KEY:
-    st.error("❌ Together API key not found. Please add it in your Hugging Face Secrets.")
+    st.error("❌ Together API key not found. Please add it under Secrets on Hugging Face.")
     st.stop()
 
-# File and input
-uploaded_file = st.file_uploader("📂 Upload your Excel file", type=["xlsx"])
+# 📂 File upload
+uploaded_file = st.file_uploader("Upload your Excel file", type=["xlsx"])
+
+# 💬 User input
 user_query = st.text_input("💬 Ask a question about your dataset")
 
-# Process on button click
-if st.button("🚀 Generate & Run SQL"):
+# 🚀 On click
+if st.button("Generate & Run SQL"):
     if uploaded_file is not None and user_query.strip():
         try:
-            # Load DataFrame
             df = pd.read_excel(uploaded_file)
             st.success("✅ File loaded successfully")
+            st.dataframe(df.head(3))
 
+            # 📋 Preview for prompt
             preview = df.head(5).to_string(index=False)
             prompt = f"""
-You are an expert data analyst. Given this sample dataset, write a Python pandas code snippet to answer the question.
+You are an AI assistant. Generate an SQL query for a pandas DataFrame named 'df' based on the user's question.
 
 Data preview:
 {preview}
 
 Question:
 {user_query}
+
+Return ONLY the SQL query. Do not include any explanations or Python code.
 """
 
-            # Make request to Together AI
-            with st.spinner("🧠 Generating response from Together AI..."):
+            # 🔗 Together API request
+            with st.spinner("🧠 Thinking..."):
                 headers = {
                     "Authorization": f"Bearer {TOGETHER_API_KEY}",
                     "Content-Type": "application/json"
                 }
                 payload = {
                     "model": "mistralai/Mixtral-8x7B-Instruct-v0.1",
-                    "max_tokens": 512,
+                    "max_tokens": 256,
                     "temperature": 0.3,
                     "messages": [
-                        {"role": "system", "content": "You are a helpful AI that writes pandas code from data questions."},
+                        {"role": "system", "content": "You generate SQL queries from natural language questions."},
                         {"role": "user", "content": prompt}
                     ]
                 }
@@ -60,21 +65,20 @@ Question:
                 )
 
                 if response.status_code != 200:
-                    raise ValueError(f"API error: {response.status_code} - {response.text}")
+                    raise ValueError(f"API Error: {response.status_code} - {response.text}")
 
-                answer = response.json()['choices'][0]['message']['content']
-                st.code(answer, language='python')
+                sql_query = response.json()['choices'][0]['message']['content'].strip()
+                st.code(sql_query, language='sql')
 
-                # Try running the generated code
+                # 🧪 Execute SQL
                 try:
-                    local_vars = {"df": df.copy()}
-                    exec(answer, {}, local_vars)
-                    if "df" in local_vars:
-                        st.dataframe(local_vars["df"])
+                    result_df = sqldf(sql_query, {"df": df})
+                    st.success("✅ SQL query executed successfully!")
+                    st.dataframe(result_df)
                 except Exception as e:
-                    st.warning(f"⚠️ Couldn't execute generated code:\n{e}")
+                    st.error(f"⚠️ Failed to run SQL query:\n{e}")
 
         except Exception as e:
-            st.error(f"❌ Error processing file: {e}")
+            st.error(f"❌ Something went wrong: {e}")
     else:
-        st.warning("⚠️ Please upload a file and enter a question.")
+        st.warning("📢 Please upload a file and enter your question.")