Update app3.py

app3.py

@@ -11,7 +11,6 @@ from langchain.chains import RetrievalQA
 from langchain.schema import Document
 import os
 import logging
-import re
 
 # Configure logging
 logging.basicConfig(level=logging.DEBUG)
@@ -53,7 +52,15 @@ else:
                 try:
                     # Load Hugging Face dataset
                     dataset = load_dataset(dataset_name, split="train", trust_remote_code=True)
-                    st.session_state.df = pd.DataFrame(dataset)
+
+                    # Convert dataset to Pandas DataFrame
+                    if isinstance(dataset, dict) or isinstance(dataset, list):
+                        st.session_state.df = pd.DataFrame(dataset)
+                    elif hasattr(dataset, 'to_pandas'):
+                        st.session_state.df = dataset.to_pandas()
+                    else:
+                        raise ValueError("Invalid input data. Cannot convert it to a DataFrame.")
+
                     st.success(f"Dataset '{dataset_name}' loaded successfully!")
                     st.dataframe(st.session_state.df.head(10))
                 except Exception as e:
@@ -85,7 +92,7 @@ else:
         try:
             # Initialize PandasAI Agent
             llm = OpenAI(api_key=pandasai_api_key, max_tokens=1500, timeout=60)
-            agent = Agent(llm)  # Corrected initialization
+            agent = Agent(df, llm=llm)
 
             # Convert the DataFrame into documents for RAG
             documents = [
@@ -108,7 +115,7 @@ else:
             )
 
             # Create tabs
-            tab1, tab2, tab3 = st.tabs(["PandasAI Analysis", "RAG Q&A", "Data Visualization"])
+            tab1, tab2, tab3 = st.tabs(["PandasAI Analysis", "RAG QA", "Data Visualization"])
 
             # Tab 1: PandasAI Analysis
             with tab1:
@@ -125,7 +132,7 @@ else:
                         st.error(f"Error from PandasAI: {e}")
                         logger.error(f"PandasAI error: {e}")
 
-            # Tab 2: RAG Q&A
+            # Tab 2: RAG QA
             with tab2:
                 st.header("Question Answering using RAG")
                 rag_question = st.text_input("Ask a question about the data (RAG):")
@@ -153,7 +160,7 @@ else:
                             viz_code = code_match.group(1)
                             logger.debug(f"Extracted visualization code: {viz_code}")
 
-                            # Modify and validate code for Plotly
+                            # Safeguard: Modify and validate code for Plotly
                             viz_code = viz_code.replace('plt.', 'px.')
                             exec(viz_code)  # Execute the visualization code
                             st.plotly_chart(fig)