Update app.py

app.py

@@ -20,10 +20,10 @@ from langchain_community.utilities.sql_database import SQLDatabase
 from datasets import load_dataset
 import tempfile
 
-# Setup API key
+# Setup API Key
 os.environ["GROQ_API_KEY"] = st.secrets.get("GROQ_API_KEY", "")
 
-# Callback handler for logging
+# LLM Logging
 class LLMCallbackHandler(BaseCallbackHandler):
     def __init__(self, log_path: Path):
         self.log_path = log_path
@@ -44,33 +44,36 @@ llm = ChatGroq(
     callbacks=[LLMCallbackHandler(Path("prompts.jsonl"))],
 )
 
-st.title("SQL-RAG using CrewAI 🚀")
+st.title("SQL-RAG Using CrewAI 🚀")
 st.write("Analyze and summarize data using natural language queries with SQL-based retrieval.")
 
-# File upload or Hugging Face dataset input
-option = st.radio("Choose your input method:", ["Upload a CSV file", "Enter Hugging Face dataset name"])
+# Primary Option: Hugging Face Dataset
+st.subheader("Option 1: Use a Hugging Face Dataset")
+default_dataset = "Einstellung/demo-salaries"
+dataset_name = st.text_input("Enter Hugging Face dataset name:", value=default_dataset)
 
-if option == "Upload a CSV file":
-    uploaded_file = st.file_uploader("Upload your dataset (CSV format)", type=["csv"])
-    if uploaded_file:
-        df = pd.read_csv(uploaded_file)
-        st.success("File uploaded successfully!")
-else:
-    dataset_name = st.text_input("Enter Hugging Face dataset name:", placeholder="e.g., imdb, ag_news")
-    if dataset_name:
-        try:
+df = None
+if dataset_name:
+    try:
+        with st.spinner("Loading Hugging Face dataset..."):
             dataset = load_dataset(dataset_name, split="train")
             df = pd.DataFrame(dataset)
             st.success(f"Dataset '{dataset_name}' loaded successfully!")
-        except Exception as e:
-            st.error(f"Error loading Hugging Face dataset: {e}")
-            df = None
-
-if 'df' in locals() and not df.empty:
-    st.write("### Dataset Preview:")
-    st.dataframe(df.head())
+            st.dataframe(df.head())
+    except Exception as e:
+        st.error(f"Error loading Hugging Face dataset: {e}")
+
+# Secondary Option: File Upload
+st.subheader("Option 2: Upload Your CSV File")
+uploaded_file = st.file_uploader("Upload your dataset (CSV format):", type=["csv"])
+if uploaded_file and df is None:
+    with st.spinner("Loading uploaded file..."):
+        df = pd.read_csv(uploaded_file)
+        st.success("File uploaded successfully!")
+        st.dataframe(df.head())
 
-    # Create a temporary SQLite database
+if df is not None:
+    # Create SQLite database
     temp_dir = tempfile.TemporaryDirectory()
     db_path = os.path.join(temp_dir.name, "data.db")
     connection = sqlite3.connect(db_path)
@@ -146,7 +149,7 @@ if 'df' in locals() and not df.empty:
         memory=False,
     )
 
-    query = st.text_input("Enter your query:", placeholder="e.g., 'What are the top 5 highest salaries?'")
+    query = st.text_input("Enter your query:", placeholder="e.g., 'What is the average salary by experience level?'")
     if query:
         with st.spinner("Processing your query..."):
             inputs = {"query": query}
@@ -156,4 +159,4 @@ if 'df' in locals() and not df.empty:
 
     temp_dir.cleanup()
 else:
-    st.warning("Please upload a valid file or provide a correct Hugging Face dataset name.")
+    st.warning("Please load a Hugging Face dataset or upload a CSV file to proceed.")