Update app.py

app.py

@@ -1,24 +1,32 @@
 import streamlit as st
 import pandas as pd
 import os
-from pandasai import SmartDataframe
+from pandasai import SmartDataframe, PandasAI
 from pandasai.llm import OpenAI
 import tempfile
 import matplotlib.pyplot as plt
 from datasets import load_dataset
-from langchain_groq import ChatGroq
-from langchain_openai import ChatOpenAI
 import time
 
+# Set Streamlit page config FIRST
+st.set_page_config(layout='wide')
 
+# Set API key
 openai_api_key = os.getenv("OPENAI_API_KEY")
 
+# Define the LLM
+llm = OpenAI(api_token=openai_api_key)
+
+# Chat with CSV
+def chat_with_csv(df, prompt):
+    pandas_ai = PandasAI(llm)
+    result = pandas_ai.run(df, prompt=prompt)
+    return result
+
 # Dataset loading without caching to support progress bar
 def load_huggingface_dataset(dataset_name):
-    # Initialize progress bar
     progress_bar = st.progress(0)
     try:
-        # Incrementally update progress
         progress_bar.progress(10)
         dataset = load_dataset(dataset_name, name="sample", split="train", trust_remote_code=True, uniform_split=True)
         progress_bar.progress(50)
@@ -26,25 +34,24 @@ def load_huggingface_dataset(dataset_name):
             df = dataset.to_pandas()
         else:
             df = pd.DataFrame(dataset)
-        progress_bar.progress(100)  # Final update to 100%
+        progress_bar.progress(100)
         return df
     except Exception as e:
-        progress_bar.progress(0)  # Reset progress bar on failure
+        progress_bar.progress(0)
         raise e
 
+# Load CSV file
 def load_uploaded_csv(uploaded_file):
-    # Initialize progress bar
     progress_bar = st.progress(0)
     try:
-        # Simulate progress
         progress_bar.progress(10)
-        time.sleep(1)  # Simulate file processing delay
+        time.sleep(1)
         progress_bar.progress(50)
         df = pd.read_csv(uploaded_file)
-        progress_bar.progress(100)  # Final update
+        progress_bar.progress(100)
         return df
     except Exception as e:
-        progress_bar.progress(0)  # Reset progress bar on failure
+        progress_bar.progress(0)
         raise e
 
 # Dataset selection logic
@@ -54,7 +61,6 @@ def load_dataset_into_session():
         ["Use Repo Directory Dataset", "Use Hugging Face Dataset", "Upload CSV File"], index=1, horizontal=True
     )
 
-    # Option 1: Load dataset from the repo directory
     if input_option == "Use Repo Directory Dataset":
         file_path = "./source/test.csv"
         if st.button("Load Dataset"):
@@ -65,11 +71,8 @@ def load_dataset_into_session():
             except Exception as e:
                 st.error(f"Error loading dataset from the repo directory: {e}")
 
-    # Option 2: Load dataset from Hugging Face
     elif input_option == "Use Hugging Face Dataset":
-        dataset_name = st.text_input(
-            "Enter Hugging Face Dataset Name:", value="HUPD/hupd"
-        )
+        dataset_name = st.text_input("Enter Hugging Face Dataset Name:", value="HUPD/hupd")
         if st.button("Load Dataset"):
             try:
                 st.session_state.df = load_huggingface_dataset(dataset_name)
@@ -77,7 +80,6 @@ def load_dataset_into_session():
             except Exception as e:
                 st.error(f"Error loading Hugging Face dataset: {e}")
 
-    # Option 3: Upload CSV File
     elif input_option == "Upload CSV File":
         uploaded_file = st.file_uploader("Upload a CSV File:", type=["csv"])
         if uploaded_file:
@@ -87,39 +89,23 @@ def load_dataset_into_session():
             except Exception as e:
                 st.error(f"Error reading uploaded file: {e}")
 
-# Load dataset into session
-load_dataset_into_session()
-
-if "df" in st.session_state and llm:
-    df = st.session_state.df
-
-    # Display dataset metadata
-    st.write("### Dataset Metadata")
-    st.text(f"Number of Rows: {df.shape[0]}")
-    st.text(f"Number of Columns: {df.shape[1]}")
-    st.text(f"Column Names: {', '.join(df.columns)}")
-
-    # Display dataset preview
-    st.write("### Dataset Preview")
-    num_rows = st.slider("Select number of rows to display:", min_value=5, max_value=50, value=10)
-    st.dataframe(df.head(num_rows))
-
-
 # Streamlit app main
-st.set_page_config(layout='wide')
 st.title("ChatCSV powered by LLM")
 
+if "df" not in st.session_state:
+    st.session_state.df = pd.DataFrame()
 
 st.header("Load Your Dataset")
 load_dataset_into_session()
 
-if not st.session_state.df.empty:
+if "df" in st.session_state and not st.session_state.df.empty:
     st.subheader("Dataset Preview")
-    st.dataframe(st.session_state.df, use_container_width=True)
+    num_rows = st.slider("Select number of rows to display:", min_value=5, max_value=50, value=10)
+    st.dataframe(st.session_state.df.head(num_rows))
 
     st.subheader("Chat with Your Dataset")
     user_query = st.text_area("Enter your query:")
-    
+
     if st.button("Run Query"):
         if user_query.strip():
             with st.spinner("Processing your query..."):