Update app.py

app.py

@@ -87,17 +87,17 @@ if st.session_state.df is not None:
 
     @tool("tables_schema")
     def tables_schema(tables: str) -> str:
-        """Get schema and sample rows for given tables."""
+        """Get the schema and sample rows for the specified tables."""
         return InfoSQLDatabaseTool(db=db).invoke(tables)
 
     @tool("execute_sql")
     def execute_sql(sql_query: str) -> str:
-        """Execute a SQL query against the database."""
+        """Execute a SQL query against the database and return the results."""
         return QuerySQLDataBaseTool(db=db).invoke(sql_query)
 
     @tool("check_sql")
     def check_sql(sql_query: str) -> str:
-        """Check the validity of a SQL query."""
+        """Validate the SQL query syntax and structure before execution."""
         return QuerySQLCheckerTool(db=db, llm=llm).invoke({"query": sql_query})
 
     sql_dev = Agent(
@@ -149,3 +149,34 @@ if st.session_state.df is not None:
         verbose=True,
     )
 
+    # UI: Tabs for Query Results and General Insights
+    tab1, tab2 = st.tabs(["🔍 Query Insights + Viz", "📊 Full Data Viz"])
+
+    with tab1:
+        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 query..."):
+                inputs = {"query": query}
+                result = crew.kickoff(inputs=inputs)
+                st.markdown("### Analysis Report:")
+                st.markdown(result)
+
+                # Query-Specific Visualization
+                if "salary" in query.lower():
+                    fig = px.box(st.session_state.df, x="job_title", y="salary_in_usd", title="Salary Distribution by Job Title")
+                    st.plotly_chart(fig)
+
+    with tab2:
+        st.subheader("📊 Comprehensive Data Visualizations")
+        
+        fig1 = px.histogram(st.session_state.df, x="job_title", title="Job Title Frequency")
+        st.plotly_chart(fig1)
+
+        fig2 = px.bar(st.session_state.df.groupby("experience_level")["salary_in_usd"].mean().reset_index(),
+                      x="experience_level", y="salary_in_usd", title="Average Salary by Experience Level")
+        st.plotly_chart(fig2)
+
+    temp_dir.cleanup()
+else:
+    st.info("Please load a dataset to proceed.")
+