Update interim.py

interim.py

@@ -20,10 +20,9 @@ from langchain_community.utilities.sql_database import SQLDatabase
 from datasets import load_dataset
 import tempfile
 
-# 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
@@ -37,103 +36,100 @@ 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")
 
-# LLM Setup
 llm = ChatGroq(
     temperature=0,
     model_name="mixtral-8x7b-32768",
     callbacks=[LLMCallbackHandler(Path("prompts.jsonl"))],
 )
 
-st.title("SQL-RAG using CrewAI 🚀")
-st.write("Analyze and summarize data using natural language queries with SQL-based retrieval.")
+st.title("SQL-RAG Using CrewAI 🚀")
+st.write("Analyze datasets using natural language queries powered by SQL and CrewAI.")
 
-# File upload or Hugging Face dataset input
-option = st.radio("Choose your input method:", ["Upload a CSV file", "Enter Hugging Face dataset name"])
+# Data Input Options
+input_option = st.radio("Select Dataset Input:", ["Use Hugging Face Dataset", "Upload CSV File"])
+df = None
 
-if option == "Upload a CSV file":
-    uploaded_file = st.file_uploader("Upload your dataset (CSV format)", type=["csv"])
+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 Hugging Face dataset..."):
+                dataset = load_dataset(dataset_name, split="train")
+                df = pd.DataFrame(dataset)
+                st.success(f"Dataset '{dataset_name}' loaded successfully!")
+                st.dataframe(df.head())
+        except Exception as e:
+            st.error(f"Error loading dataset: {e}")
+else:
+    uploaded_file = st.file_uploader("Upload CSV File:", 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:
-            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())
 
-    # Create a temporary SQLite database
+# SQL-RAG and Query Workflow
+if 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("data_table", connection, if_exists="replace", index=False)
+    df.to_sql("salaries", connection, if_exists="replace", index=False)
     db = SQLDatabase.from_uri(f"sqlite:///{db_path}")
 
-    # Tools
     @tool("list_tables")
     def list_tables() -> str:
+        """List all tables in the database."""
         return ListSQLDatabaseTool(db=db).invoke("")
 
     @tool("tables_schema")
     def tables_schema(tables: str) -> str:
+        """Return schema and example rows for given tables."""
         return InfoSQLDatabaseTool(db=db).invoke(tables)
 
     @tool("execute_sql")
     def execute_sql(sql_query: str) -> str:
+        """Execute a SQL query and return results."""
         return QuerySQLDataBaseTool(db=db).invoke(sql_query)
 
     @tool("check_sql")
     def check_sql(sql_query: str) -> str:
+        """Check SQL query validity."""
         return QuerySQLCheckerTool(db=db, llm=llm).invoke({"query": sql_query})
 
-    # Agents
     sql_dev = Agent(
-        role="Database Developer",
-        goal="Extract data from the database.",
+        role="Senior Database Developer",
+        goal="Construct and execute SQL queries.",
         llm=llm,
         tools=[list_tables, tables_schema, execute_sql, check_sql],
-        allow_delegation=False,
     )
 
     data_analyst = Agent(
-        role="Data Analyst",
-        goal="Analyze and provide insights.",
+        role="Senior Data Analyst",
+        goal="Analyze the data returned from SQL queries.",
         llm=llm,
-        allow_delegation=False,
     )
 
     report_writer = Agent(
-        role="Report Editor",
-        goal="Summarize the analysis.",
+        role="Senior Report Editor",
+        goal="Summarize the analysis into a short report.",
         llm=llm,
-        allow_delegation=False,
     )
 
-    # Tasks
     extract_data = Task(
-        description="Extract data required for the query: {query}.",
-        expected_output="Database result for the query",
+        description="Extract data for the query: {query}.",
+        expected_output="Database query results.",
         agent=sql_dev,
     )
 
     analyze_data = Task(
-        description="Analyze the data for: {query}.",
-        expected_output="Detailed analysis text",
+        description="Analyze the query results for: {query}.",
+        expected_output="Detailed analysis report.",
         agent=data_analyst,
         context=[extract_data],
     )
 
     write_report = Task(
-        description="Summarize the analysis into a short report.",
-        expected_output="Markdown report",
+        description="Summarize the analysis into a brief executive summary.",
+        expected_output="Markdown report.",
         agent=report_writer,
         context=[analyze_data],
     )
@@ -143,12 +139,11 @@ if 'df' in locals() and not df.empty:
         tasks=[extract_data, analyze_data, write_report],
         process=Process.sequential,
         verbose=2,
-        memory=False,
     )
 
-    query = st.text_input("Enter your query:", placeholder="e.g., 'What are the top 5 highest salaries?'")
-    if query:
-        with st.spinner("Processing your query..."):
+    query = st.text_area("Enter Query:", placeholder="e.g., 'What is the average salary by experience level?'")
+    if st.button("Submit Query"):
+        with st.spinner("Processing your query with CrewAI..."):
             inputs = {"query": query}
             result = crew.kickoff(inputs=inputs)
             st.markdown("### Analysis Report:")
@@ -156,4 +151,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.info("Load a dataset to proceed.")