Update src/streamlit_app.py

src/streamlit_app.py

@@ -12,7 +12,7 @@ import requests
 from sqlalchemy import create_engine, text, inspect
 
 # --- Get HF Token ---
-HF_TOKEN = os.environ["HF_TOKEN"]  # will raise KeyError if not set
+HF_TOKEN = os.environ.get("HF_TOKEN", "")  # Safely get token, fallback if missing
 
 # --- Helper: Call Mistral Model ---
 def mistral_call(schema=None, question="no questions were asked", hf_token=HF_TOKEN, model_id="mistralai/Mistral-7B-Instruct-v0.3"):
@@ -153,32 +153,45 @@ option = st.sidebar.radio("Select Feature", ["📊 Data Visualization", "🧠 SQ
 if option == "📊 Data Visualization":
     uploaded_file = st.file_uploader("Upload your CSV", type="csv")
     if uploaded_file:
-        df = pd.read_csv(uploaded_file)
-        df.columns = df.columns.str.strip()
-        st.dataframe(df.head())
-        with st.spinner("Getting chart suggestion..."):
-            suggestion = get_visualization_suggestion(df)
-        if suggestion:
-            x_col = suggestion.get("x", "").strip()
-            y_col = suggestion.get("y", [])
-            y_col = [y_col] if isinstance(y_col, str) else y_col
-            chart = suggestion.get("chart_type")
-            if x_col in df.columns and all(y in df.columns for y in y_col):
-                fig = None
-                if chart == "bar":
-                    fig = px.bar(df, x=x_col, y=y_col)
-                elif chart == "line":
-                    fig = px.line(df, x=x_col, y=y_col)
-                elif chart == "scatter":
-                    fig = px.scatter(df, x=x_col, y=y_col)
-                elif chart == "pie" and len(y_col) == 1:
-                    fig = px.pie(df, names=x_col, values=y_col[0])
-                if fig:
-                    st.plotly_chart(fig)
+        try:
+            content = uploaded_file.getvalue().decode("utf-8")
+            df = pd.read_csv(io.StringIO(content))
+            df.columns = df.columns.str.strip().str.replace(" ", "_")
+            st.write("CSV Preview")
+            st.dataframe(df.head())
+            st.write("Shape:", df.shape)
+
+            with st.spinner("Getting chart suggestion..."):
+                suggestion = get_visualization_suggestion(df)
+
+            st.write("Model suggestion:")
+            st.code(suggestion)
+
+            if suggestion:
+                x_col = suggestion.get("x", "").strip()
+                y_col = suggestion.get("y", [])
+                y_col = [y_col] if isinstance(y_col, str) else y_col
+                chart = suggestion.get("chart_type")
+                if x_col in df.columns and all(y in df.columns for y in y_col):
+                    fig = None
+                    if chart == "bar":
+                        fig = px.bar(df, x=x_col, y=y_col)
+                    elif chart == "line":
+                        fig = px.line(df, x=x_col, y=y_col)
+                    elif chart == "scatter":
+                        fig = px.scatter(df, x=x_col, y=y_col)
+                    elif chart == "pie" and len(y_col) == 1:
+                        fig = px.pie(df, names=x_col, values=y_col[0])
+                    if fig:
+                        st.plotly_chart(fig)
+                    else:
+                        st.error("Unsupported chart type.")
                 else:
-                    st.error("Unsupported chart type.")
+                    st.error("⚠️ Column suggestion doesn't match your CSV.")
             else:
-                st.error("Invalid column suggestion from model.")
+                st.error("❌ No valid visualization suggestion returned.")
+        except Exception as e:
+            st.error(f"❌ Error reading CSV: {e}")
 
 elif option == "🧠 SQL Query Generator":
     user_input = st.text_area("Describe your SQL query in plain English:")