Update app.py

app.py

@@ -20,7 +20,7 @@ from langchain_community.utilities.sql_database import SQLDatabase
 from datasets import load_dataset
 import tempfile
 
-# Environment setup
+# API Key
 os.environ["GROQ_API_KEY"] = st.secrets.get("GROQ_API_KEY", "")
 
 # LLM Callback Logger
@@ -37,17 +37,18 @@ 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 the LLM
+# 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.")
 
-# Input Options
+# Dataset Input
 input_option = st.radio("Select Dataset Input:", ["Use Hugging Face Dataset", "Upload CSV File"])
 df = None
 
@@ -55,13 +56,13 @@ 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..."):
+            with st.spinner("Loading 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}")
+            st.error(f"Error: {e}")
 else:
     uploaded_file = st.file_uploader("Upload CSV File:", type=["csv"])
     if uploaded_file:
@@ -77,79 +78,66 @@ if df is not None:
     df.to_sql("salaries", connection, if_exists="replace", index=False)
     db = SQLDatabase.from_uri(f"sqlite:///{db_path}")
 
-    # Tools with proper docstrings
     @tool("list_tables")
     def list_tables() -> str:
-        """List all tables in the SQLite database."""
+        """List all tables in the database."""
         return ListSQLDatabaseTool(db=db).invoke("")
 
     @tool("tables_schema")
     def tables_schema(tables: str) -> str:
-        """
-        Get the schema and sample rows for specific tables in the database.
-
-        Input: Comma-separated table names.
-        Example: 'salaries'
-        """
+        """Get schema and sample rows for given tables."""
         return InfoSQLDatabaseTool(db=db).invoke(tables)
 
     @tool("execute_sql")
     def execute_sql(sql_query: str) -> str:
-        """
-        Execute a valid SQL query on the database and return the results.
-
-        Input: A SQL query string.
-        Example: 'SELECT * FROM salaries LIMIT 5;'
-        """
+        """Execute a SQL query against the database."""
         return QuerySQLDataBaseTool(db=db).invoke(sql_query)
 
     @tool("check_sql")
     def check_sql(sql_query: str) -> str:
-        """
-        Check the validity of a SQL query before execution.
-
-        Input: A SQL query string.
-        Example: 'SELECT salary FROM salaries WHERE salary > 10000;'
-        """
+        """Check the validity of a SQL query."""
         return QuerySQLCheckerTool(db=db, llm=llm).invoke({"query": sql_query})
 
     # Agents
     sql_dev = Agent(
-        role="Database Developer",
-        goal="Extract relevant data by executing SQL queries.",
+        role="Senior Database Developer",
+        goal="Extract data using optimized SQL queries.",
+        backstory="An expert in writing optimized SQL queries for complex databases.",
         llm=llm,
         tools=[list_tables, tables_schema, execute_sql, check_sql],
     )
 
     data_analyst = Agent(
-        role="Data Analyst",
-        goal="Analyze the extracted data and generate detailed insights.",
+        role="Senior Data Analyst",
+        goal="Analyze the data and produce insights.",
+        backstory="A seasoned analyst who identifies trends and patterns in datasets.",
         llm=llm,
     )
 
     report_writer = Agent(
-        role="Report Writer",
-        goal="Summarize the analysis into an executive report.",
+        role="Technical Report Writer",
+        goal="Summarize the insights into a clear report.",
+        backstory="An expert in summarizing data insights into readable reports.",
         llm=llm,
     )
 
     # Tasks
     extract_data = Task(
-        description="Extract data for the query: {query}.",
-        expected_output="Database query results.",
+        description="Extract data based on the query: {query}.",
+        expected_output="Database results matching the query.",
         agent=sql_dev,
     )
 
     analyze_data = Task(
-        description="Analyze the query results for: {query}.",
-        expected_output="Analysis report.",
+        description="Analyze the extracted data for query: {query}.",
+        expected_output="Analysis text summarizing findings.",
         agent=data_analyst,
         context=[extract_data],
     )
 
     write_report = Task(
-        description="Summarize the analysis into an executive summary.",
-        expected_output="Markdown-formatted report.",
+        description="Summarize the analysis into an executive report.",
+        expected_output="Markdown report of insights.",
         agent=report_writer,
         context=[analyze_data],
     )
@@ -161,9 +149,9 @@ if df is not None:
         verbose=True,
     )
 
-    query = st.text_area("Enter Query:", placeholder="e.g., 'What is the average salary by experience level?'")
+    query = st.text_area("Enter Query:", placeholder="e.g., 'What is the average salary for senior employees?'")
     if st.button("Submit Query"):
-        with st.spinner("Processing your query with CrewAI..."):
+        with st.spinner("Processing query..."):
             inputs = {"query": query}
             result = crew.kickoff(inputs=inputs)
             st.markdown("### Analysis Report:")
@@ -171,4 +159,4 @@ if df is not None:
 
     temp_dir.cleanup()
 else:
-    st.info("Load a dataset to proceed.")
+    st.info("Please load a dataset to proceed.")