Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,76 +1,66 @@
 import streamlit as st
 import pandas as pd
-import duckdb
-import requests
-import re
 import os
+from groq import Groq
 
-st.set_page_config(page_title="🧠 SQL Chatbot with Groq", layout="centered")
+# Set page configuration
+st.set_page_config(page_title="📊 Excel SQL Assistant", layout="centered")
+st.title("🧠 Excel to SQL with GROQ + Streamlit")
 
-# 🔐 Load Groq API Key from Hugging Face secret or env var
-GROQ_API_KEY = os.getenv("GROQ_API_KEY")
+# Load Groq API key from Streamlit secrets
+GROQ_API_KEY = st.secrets.get("GROQ_API_KEY", None)
 
+# Fallback warning if API key not found
 if not GROQ_API_KEY:
-    st.error("❌ GROQ_API_KEY not found. Please set it in Streamlit Secrets or environment.")
+    st.error("❌ GROQ API key not found. Please set it in your Hugging Face `secrets` tab.")
     st.stop()
 
-# 🧠 Generate SQL using Groq API
-def generate_sql(prompt, df):
-    schema = ", ".join([f"{col} ({dtype})" for col, dtype in df.dtypes.items()])
-    full_prompt = f"""
-You are a SQL expert. The table 'df' has the following columns and types:
-{schema}
+# Initialize Groq client
+client = Groq(api_key=GROQ_API_KEY)
 
-User question: "{prompt}"
+# Upload Excel file
+uploaded_file = st.file_uploader("📂 Upload your Excel file", type=["xlsx"])
+st.write("📁 Uploaded file object:", uploaded_file)  # Debugging aid
 
-Write a valid SQL query using the 'df' table. Return only the SQL.
-"""
+# Text input for user query
+user_query = st.text_input("💬 Ask a question about your data")
 
-    headers = {
-        "Authorization": f"Bearer {GROQ_API_KEY}",
-        "Content-Type": "application/json"
-    }
-    payload = {
-        "model": "llama3-70b-8192",
-        "messages": [{"role": "user", "content": full_prompt}],
-        "temperature": 0.2,
-        "max_tokens": 300
-    }
+# Generate button
+if st.button("🚀 Generate & Run SQL"):
+    if uploaded_file is not None and user_query.strip() != "":
+        try:
+            df = pd.read_excel(uploaded_file)
+            st.success("✅ File successfully loaded!")
 
-    url = "https://api.groq.com/openai/v1/chat/completions"
-    response = requests.post(url, headers=headers, json=payload)
-    response.raise_for_status()
-    result = response.json()
-    return result["choices"][0]["message"]["content"].strip()
+            # Create schema-like string from the DataFrame
+            preview = df.head(5).to_string(index=False)
+            schema = f"Here are the first few rows of the dataset:\n\n{preview}"
 
-# 🧽 Clean SQL
-def clean_sql(sql, df_columns):
-    sql = sql.replace("`", '"')
-    for col in df_columns:
-        if " " in col and f'"{col}"' not in sql:
-            pattern = r'\b' + re.escape(col) + r'\b'
-            sql = re.sub(pattern, f'"{col}"', sql)
-    return sql
+            # Prompt to Groq
+            full_prompt = f"""You are a data expert. Write a Pandas code snippet to answer the following question based on the data:\n\n{schema}\n\nQuestion: {user_query}"""
 
-# 🖥️ UI
-st.title("📊 Excel SQL Chatbot (Groq + Streamlit)")
-uploaded_file = st.file_uploader("📂 Upload your Excel file", type=["xlsx"])
-user_prompt = st.text_input("🧠 Ask a question about your data")
+            with st.spinner("⏳ Generating SQL-like response using Groq..."):
+                response = client.chat.completions.create(
+                    model="mixtral-8x7b-32768",
+                    messages=[
+                        {"role": "system", "content": "You are an expert data assistant that writes pandas code based on user questions."},
+                        {"role": "user", "content": full_prompt}
+                    ]
+                )
 
-if st.button("🚀 Generate & Run SQL"):
-    if uploaded_file and user_prompt:
-        try:
-            df = pd.read_excel(uploaded_file)
-            sql = generate_sql(user_prompt, df)
-            cleaned_sql = clean_sql(sql, df.columns)
-            result = duckdb.query(cleaned_sql).to_df()
+            answer = response.choices[0].message.content
+            st.code(answer, language="python")
 
-            st.markdown(f"### 🧾 Generated SQL")
-            st.code(sql, language="sql")
-            st.markdown("### 📈 Query Results")
-            st.dataframe(result)
+            # Optionally: try to run it
+            try:
+                local_vars = {"df": df.copy()}
+                exec(answer, {}, local_vars)
+                if 'df' in local_vars:
+                    st.dataframe(local_vars['df'])
+            except Exception as e:
+                st.warning(f"⚠️ Could not execute the code:\n\n{e}")
 
         except Exception as e:
-            st.error(f"❌ Error: {str(e)}")
+            st.error(f"❌ Error processing file: {e}")
     else:
-        st.warning("⚠️ Please upload a file and enter a question.")
+        st.warning("⚠️ Please upload a file and enter a question.")