Update app.py

app.py

@@ -115,26 +115,32 @@ def generate_visualization(suggestion, df):
     y_axis = suggestion.get("y_axis")
     group_by = suggestion.get("group_by")
 
-    # Ensure required inputs are available
+    # Dynamically determine the best Y-axis if GPT-4o doesn't suggest one
+    if not y_axis:
+        numeric_columns = df.select_dtypes(include='number').columns.tolist()
+        
+        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
     if not x_axis or not y_axis:
-        st.warning("⚠️ GPT-4o did not provide enough information for the visualization.")
+        st.warning("⚠️ Unable to determine relevant columns for visualization.")
         return None
 
     # Dynamically select the Plotly function
     plotly_function = getattr(px, chart_type, None)
 
-    # Handle unsupported chart types gracefully
     if not plotly_function:
         st.warning(f"⚠️ Unsupported chart type '{chart_type}' suggested by GPT-4o.")
         return None
 
-    # Prepare dynamic parameters for Plotly function
-    plot_args = {
-        "data_frame": df,
-        "x": x_axis,
-        "y": y_axis,
-    }
-    if group_by:
+    # 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:
@@ -145,12 +151,17 @@ def generate_visualization(suggestion, df):
             xaxis_title=x_axis.replace('_', ' ').title(),
             yaxis_title=y_axis.replace('_', ' ').title(),
         )
+
+        # Apply statistics intelligently
+        fig = add_stats_to_figure(fig, df, y_axis, chart_type)
+
         return fig
 
     except Exception as e:
         st.error(f"⚠️ Failed to generate visualization: {e}")
         return None
 
+
 def add_stats_to_figure(fig, df, y_axis, chart_type):
     # Calculate statistics
     min_val = df[y_axis].min()