Update app.py

app.py

@@ -23,7 +23,7 @@ import tempfile
 # API Key
 os.environ["GROQ_API_KEY"] = st.secrets.get("GROQ_API_KEY", "")
 
-# LLM Callback Logger
+# Initialize LLM
 class LLMCallbackHandler(BaseCallbackHandler):
     def __init__(self, log_path: Path):
         self.log_path = log_path
@@ -37,45 +37,45 @@ class LLMCallbackHandler(BaseCallbackHandler):
         with self.log_path.open("a", encoding="utf-8") as file:
             file.write(json.dumps({"event": "llm_end", "text": generation, "timestamp": datetime.now().isoformat()}) + "\n")
 
-# Initialize LLM
 llm = ChatGroq(
     temperature=0,
     model_name="mixtral-8x7b-32768",
     callbacks=[LLMCallbackHandler(Path("prompts.jsonl"))],
 )
 
-# Streamlit UI
 st.title("SQL-RAG Using CrewAI 🚀")
 st.write("Analyze datasets using natural language queries powered by SQL and CrewAI.")
 
+# Initialize session state for data persistence
+if "df" not in st.session_state:
+    st.session_state.df = None
+
 # Dataset Input
 input_option = st.radio("Select Dataset Input:", ["Use Hugging Face Dataset", "Upload CSV File"])
-df = None
-
 if input_option == "Use Hugging Face Dataset":
     dataset_name = st.text_input("Enter Hugging Face Dataset Name:", value="Einstellung/demo-salaries")
     if st.button("Load Dataset"):
         try:
             with st.spinner("Loading dataset..."):
                 dataset = load_dataset(dataset_name, split="train")
-                df = pd.DataFrame(dataset)
+                st.session_state.df = pd.DataFrame(dataset)
                 st.success(f"Dataset '{dataset_name}' loaded successfully!")
-                st.dataframe(df.head())
+                st.dataframe(st.session_state.df.head())
         except Exception as e:
             st.error(f"Error: {e}")
-else:
+elif input_option == "Upload CSV File":
     uploaded_file = st.file_uploader("Upload CSV File:", type=["csv"])
     if uploaded_file:
-        df = pd.read_csv(uploaded_file)
+        st.session_state.df = pd.read_csv(uploaded_file)
         st.success("File uploaded successfully!")
-        st.dataframe(df.head())
+        st.dataframe(st.session_state.df.head())
 
 # SQL-RAG Analysis
-if df is not None:
+if st.session_state.df is not None:
     temp_dir = tempfile.TemporaryDirectory()
     db_path = os.path.join(temp_dir.name, "data.db")
     connection = sqlite3.connect(db_path)
-    df.to_sql("salaries", connection, if_exists="replace", index=False)
+    st.session_state.df.to_sql("salaries", connection, if_exists="replace", index=False)
     db = SQLDatabase.from_uri(f"sqlite:///{db_path}")
 
     @tool("list_tables")
@@ -98,7 +98,6 @@ if df is not None:
         """Check the validity of a SQL query."""
         return QuerySQLCheckerTool(db=db, llm=llm).invoke({"query": sql_query})
 
-    # Agents
     sql_dev = Agent(
         role="Senior Database Developer",
         goal="Extract data using optimized SQL queries.",
@@ -121,7 +120,6 @@ if df is not None:
         llm=llm,
     )
 
-    # Tasks
     extract_data = Task(
         description="Extract data based on the query: {query}.",
         expected_output="Database results matching the query.",