Update app.py

app.py

@@ -109,14 +109,23 @@ def ask_gpt4o_for_visualization(query, df, llm):
         return None
 
 def add_stats_to_figure(fig, df, y_axis, chart_type):
-    # Calculate statistics
+    """
+    Add relevant statistical annotations to the visualization 
+    based on the chart type.
+    """
+    # Check if the y-axis column is numeric
+    if not pd.api.types.is_numeric_dtype(df[y_axis]):
+        st.warning(f"⚠️ Cannot compute statistics for non-numeric column: {y_axis}")
+        return fig
+
+    # Compute statistics for numeric data
     min_val = df[y_axis].min()
     max_val = df[y_axis].max()
     avg_val = df[y_axis].mean()
     median_val = df[y_axis].median()
     std_dev_val = df[y_axis].std()
 
-    # Stats summary text
+    # Format the stats for display
     stats_text = (
         f"📊 **Statistics**\n\n"
         f"- **Min:** ${min_val:,.2f}\n"
@@ -126,42 +135,59 @@ def add_stats_to_figure(fig, df, y_axis, chart_type):
         f"- **Std Dev:** ${std_dev_val:,.2f}"
     )
 
-    # Charts suitable for stats annotations
-    if chart_type in ["bar", "line", "scatter"]:
-        # Add annotation box
+    # Apply stats only to relevant chart types
+    if chart_type in ["bar", "line"]:
+        # Add annotation box for bar and line charts
         fig.add_annotation(
             text=stats_text,
             xref="paper", yref="paper",
-            x=1.05, y=1,
+            x=1.02, y=1,
             showarrow=False,
             align="left",
             font=dict(size=12, color="black"),
-            bordercolor="black",
+            bordercolor="gray",
             borderwidth=1,
-            bgcolor="rgba(255, 255, 255, 0.8)"
+            bgcolor="rgba(255, 255, 255, 0.85)"
         )
 
-        # Add horizontal lines for min, median, avg, max
+        # Add horizontal reference lines
         fig.add_hline(y=min_val, line_dash="dot", line_color="red", annotation_text="Min", annotation_position="bottom right")
         fig.add_hline(y=median_val, line_dash="dash", line_color="orange", annotation_text="Median", annotation_position="top right")
         fig.add_hline(y=avg_val, line_dash="dashdot", line_color="green", annotation_text="Avg", annotation_position="top right")
         fig.add_hline(y=max_val, line_dash="dot", line_color="blue", annotation_text="Max", annotation_position="top right")
 
+    elif chart_type == "scatter":
+        # Add stats annotation only, no lines for scatter plots
+        fig.add_annotation(
+            text=stats_text,
+            xref="paper", yref="paper",
+            x=1.02, y=1,
+            showarrow=False,
+            align="left",
+            font=dict(size=12, color="black"),
+            bordercolor="gray",
+            borderwidth=1,
+            bgcolor="rgba(255, 255, 255, 0.85)"
+        )
+
     elif chart_type == "box":
-        # Box plots already show distribution (no extra stats needed)
+        # Box plots inherently show distribution; no extra stats needed
         pass
 
     elif chart_type == "pie":
-        # Pie charts don't need statistical overlays
-        st.info("📊 Pie charts focus on proportions. No additional stats displayed.")
-    
+        # Pie charts represent proportions, not suitable for stats
+        st.info("📊 Pie charts represent proportions. Additional stats are not applicable.")
+
+    elif chart_type == "heatmap":
+        # Heatmaps already reflect data intensity
+        st.info("📊 Heatmaps inherently reflect distribution. No additional stats added.")
+
     else:
-        st.warning(f"⚠️ No stats added for unsupported chart type: {chart_type}")
+        st.warning(f"⚠️ No statistical overlays applied for unsupported chart type: '{chart_type}'.")
 
     return fig
 
 
-
 # Dynamically generate Plotly visualizations based on GPT-4o suggestions
 def generate_visualization(suggestion, df):
     chart_type = suggestion.get("chart_type", "bar").lower()