Update app.py

app.py

@@ -194,41 +194,52 @@ def add_stats_to_figure(fig, df, y_axis, chart_type):
 
 # Dynamically generate Plotly visualizations based on GPT-4o suggestions
 def generate_visualization(suggestion, df):
+    """
+    Generate a Plotly visualization based on GPT-4o's suggestion.
+    If the Y-axis is missing, infer it intelligently.
+    """
     chart_type = suggestion.get("chart_type", "bar").lower()
     x_axis = suggestion.get("x_axis")
     y_axis = suggestion.get("y_axis")
     group_by = suggestion.get("group_by")
 
-    # Dynamically determine the best Y-axis if GPT-4o doesn't suggest one
+    # Step 1: Infer Y-axis if not provided
     if not y_axis:
         numeric_columns = df.select_dtypes(include='number').columns.tolist()
-        
+
+        # Avoid using the same column for both axes
         if x_axis in numeric_columns:
-            # Avoid using the same column for both axes
             numeric_columns.remove(x_axis)
-        
-        # Prioritize the first available numeric column for y-axis
-        y_axis = numeric_columns[0] if numeric_columns else None
 
-    # Ensure both axes are identified
+        # Smart guess: prioritize salary or relevant metrics if available
+        priority_columns = ["salary_in_usd", "income", "earnings", "revenue"]
+        for col in priority_columns:
+            if col in numeric_columns:
+                y_axis = col
+                break
+
+        # Fallback to the first numeric column if no priority columns exist
+        if not y_axis and numeric_columns:
+            y_axis = numeric_columns[0]
+
+    # Step 2: Validate axes
     if not x_axis or not y_axis:
-        st.warning("⚠️ Unable to determine relevant columns for visualization.")
+        st.warning("⚠️ Unable to determine appropriate columns for visualization.")
         return None
 
-    # Dynamically select the Plotly function
+    #  Step 3: Dynamically select the Plotly function
     plotly_function = getattr(px, chart_type, None)
-
     if not plotly_function:
         st.warning(f"⚠️ Unsupported chart type '{chart_type}' suggested by GPT-4o.")
         return None
 
-    # Prepare dynamic plot arguments
+    #  Step 4: Prepare dynamic plot arguments
     plot_args = {"data_frame": df, "x": x_axis, "y": y_axis}
     if group_by and group_by in df.columns:
         plot_args["color"] = group_by
 
     try:
-        # Generate the dynamic visualization
+        # Step 5: Generate the visualization
         fig = plotly_function(**plot_args)
         fig.update_layout(
             title=f"{chart_type.title()} Plot of {y_axis.replace('_', ' ').title()} by {x_axis.replace('_', ' ').title()}",
@@ -236,8 +247,8 @@ def generate_visualization(suggestion, df):
             yaxis_title=y_axis.replace('_', ' ').title(),
         )
 
-        # Apply statistics intelligently
-        fig = add_stats_to_figure(fig, df, y_axis, chart_type)
+        # Step 6: Apply statistics intelligently
+        fig = add_statistics_to_visualization(fig, df, y_axis, chart_type)
 
         return fig
 
@@ -245,6 +256,7 @@ def generate_visualization(suggestion, df):
         st.error(f"⚠️ Failed to generate visualization: {e}")
         return None
 
+
 def generate_multiple_visualizations(suggestions, df):
     """
     Generates one or more visualizations based on GPT-4o's suggestions.