Update app4.py

app4.py

@@ -19,6 +19,21 @@ pandasai_api_key = os.getenv("PANDASAI_API_KEY")
 if not api_key or not pandasai_api_key:
     st.warning("API keys for OpenAI or PandasAI are missing. Ensure both keys are set in environment variables.")
 
+# Add session reset button
+if st.button("Reset Session"):
+    for key in list(st.session_state.keys()):
+        del st.session_state[key]
+    st.experimental_rerun()
+
+# Function to validate and clean dataset
+def validate_and_clean_dataset(df):
+    # Rename columns for consistency
+    df.columns = [col.strip().lower().replace(" ", "_") for col in df.columns]
+    # Check for missing values
+    if df.isnull().values.any():
+        st.warning("Dataset contains missing values. Consider cleaning the data.")
+    return df
+
 # Function to load datasets into session
 def load_dataset_into_session():
     input_option = st.radio(
@@ -32,6 +47,7 @@ def load_dataset_into_session():
         if st.button("Load Dataset"):
             try:
                 st.session_state.df = pd.read_csv(file_path)
+                st.session_state.df = validate_and_clean_dataset(st.session_state.df)
                 st.success(f"File loaded successfully from '{file_path}'!")
             except Exception as e:
                 st.error(f"Error loading dataset from the repo directory: {e}")
@@ -49,6 +65,7 @@ def load_dataset_into_session():
                     st.session_state.df = dataset.to_pandas()
                 else:
                     st.session_state.df = pd.DataFrame(dataset)
+                st.session_state.df = validate_and_clean_dataset(st.session_state.df)
                 st.success(f"Hugging Face Dataset '{dataset_name}' loaded successfully!")
             except Exception as e:
                 st.error(f"Error loading Hugging Face dataset: {e}")
@@ -59,6 +76,7 @@ def load_dataset_into_session():
         if uploaded_file:
             try:
                 st.session_state.df = pd.read_csv(uploaded_file)
+                st.session_state.df = validate_and_clean_dataset(st.session_state.df)
                 st.success("File uploaded successfully!")
             except Exception as e:
                 st.error(f"Error reading uploaded file: {e}")
@@ -119,11 +137,10 @@ if "df" in st.session_state and api_key and pandasai_api_key:
             try:
                 result = agent.chat(pandas_question)
                 st.write("PandasAI Answer:", result)
+                if hasattr(agent, "last_output"):
+                    st.write("PandasAI Intermediate Output:", agent.last_output)
             except Exception as e:
                 st.error(f"PandasAI encountered an error: {str(e)}")
-            finally:
-                st.write("PandasAI intermediate output (if any):")
-                st.write(agent.last_output if hasattr(agent, "last_output") else "No intermediate output available.")
 
     with tab2:
         st.subheader("Q&A with RAG")
@@ -134,12 +151,6 @@ if "df" in st.session_state and api_key and pandasai_api_key:
                 st.write("RAG Answer:", result)
             except Exception as e:
                 st.error(f"RAG encountered an error: {str(e)}")
-            finally:
-                st.write("RAG Intermediate Status:")
-                st.write({
-                    "retriever": retriever,
-                    "qa_chain": qa_chain
-                })
 
     with tab3:
         st.subheader("Data Visualization")
@@ -147,7 +158,6 @@ if "df" in st.session_state and api_key and pandasai_api_key:
         if viz_question:
             try:
                 result = agent.chat(viz_question)
-                
                 # Extract Python code from PandasAI response
                 import re
                 code_pattern = r'```python\n(.*?)\n```'
@@ -155,21 +165,13 @@ if "df" in st.session_state and api_key and pandasai_api_key:
                 
                 if code_match:
                     viz_code = code_match.group(1)
-                    
-                    # Replace matplotlib with plotly
-                    viz_code = viz_code.replace('plt.', 'px.')
-                    viz_code = viz_code.replace('plt.show()', 'fig = px.scatter(df, x=x, y=y)')
-                    
-                    # Execute the modified code
                     exec(viz_code)
-                    st.plotly_chart(fig)
                 else:
-                    st.write("Unable to generate the graph. Please try a different query.")
+                    st.write("Unable to generate the graph. Showing fallback example.")
+                    fig = px.scatter(df, x=df.columns[0], y=df.columns[1])
+                    st.plotly_chart(fig)
             except Exception as e:
                 st.error(f"An error occurred during visualization: {str(e)}")
-            finally:
-                st.write("Visualization debug details:")
-                st.write({"viz_question": viz_question, "result": result})
 else:
     if not api_key:
         st.warning("Please set the OpenAI API key in environment variables.")