Update app.py

app.py

@@ -246,35 +246,38 @@ def add_stats_to_figure(fig, df, y_axis):
     )
     return fig
 
-# Unified visualization function with LLM fallback
+# Unified Visualization Generator with Fuzzy Matching and LLM Fallback
 def generate_visual_from_query(query, df, llm=None):
     try:
+        # Step 1: Attempt Fuzzy Matching
         matched_columns = fuzzy_match_columns(query)
 
-        # Fallback to LLM if fuzzy matching fails
+        # Step 2: Fallback to LLM if no columns are matched
         if not matched_columns and llm:
             st.info("🤖 No match found. Asking AI for suggestions...")
             suggestion = ask_llm_for_columns(query, llm, df)
             if suggestion:
                 matched_columns = [suggestion.get("x_axis"), suggestion.get("group_by")]
 
-        # Handle cases when we have columns to plot
+        # Step 3: Process Matched Columns
         if len(matched_columns) >= 2:
             x_axis, group_by = matched_columns[0], matched_columns[1]
         elif len(matched_columns) == 1:
             x_axis, group_by = matched_columns[0], None
         else:
-            st.warning("❓ No matching columns found. Please refine your query.")
+            st.warning("❓ No matching columns found. Try rephrasing your query.")
             return None
 
-        # Handle distribution queries
+        # Step 4: Visualization Generation
+
+        # Distribution Plot
         if "distribution" in query:
             fig = px.box(df, x=x_axis, y="salary_in_usd", color=group_by,
                          title=f"Salary Distribution by {x_axis.replace('_', ' ').title()}"
                                + (f" and {group_by.replace('_', ' ').title()}" if group_by else ""))
             return add_stats_to_figure(fig, df, "salary_in_usd")
 
-        # Handle average salary queries
+        # Average Salary Plot
         elif "average" in query or "mean" in query:
             grouped_df = df.groupby([x_axis] + ([group_by] if group_by else []))["salary_in_usd"].mean().reset_index()
             fig = px.bar(grouped_df, x=x_axis, y="salary_in_usd", color=group_by,
@@ -282,22 +285,23 @@ def generate_visual_from_query(query, df, llm=None):
                                + (f" and {group_by.replace('_', ' ').title()}" if group_by else ""))
             return add_stats_to_figure(fig, df, "salary_in_usd")
 
-        # Handle salary trends over time
+        # Salary Trends Over Time
         elif "trend" in query and "work_year" in df.columns:
             grouped_df = df.groupby(["work_year", x_axis])["salary_in_usd"].mean().reset_index()
             fig = px.line(grouped_df, x="work_year", y="salary_in_usd", color=x_axis,
                           title=f"Salary Trend Over Years by {x_axis.replace('_', ' ').title()}")
             return add_stats_to_figure(fig, df, "salary_in_usd")
 
-        # Handle remote work impact
+        # Remote Work Impact
         elif "remote" in query:
             grouped_df = df.groupby(["remote_ratio"] + ([group_by] if group_by else []))["salary_in_usd"].mean().reset_index()
             fig = px.bar(grouped_df, x="remote_ratio", y="salary_in_usd", color=group_by,
                          title="Remote Work Impact on Salary")
             return add_stats_to_figure(fig, df, "salary_in_usd")
 
+        # No Specific Match
         else:
-            st.warning("⚠️ No suitable visualization generated. Please refine your query.")
+            st.warning("⚠️ No suitable visualization to display!")
             return None
 
     except Exception as e:
@@ -305,7 +309,6 @@ def generate_visual_from_query(query, df, llm=None):
         return None
 
 
-
 # SQL-RAG Analysis
 if st.session_state.df is not None:
     temp_dir = tempfile.TemporaryDirectory()