Update ai_agent.py

ai_agent.py

@@ -168,6 +168,26 @@ def create_chart(df, selected_columns=None, chart_type="bar", title=None):
     except:
         return None
 
+def find_column(df, col_name):
+    """Find a column in DataFrame with fuzzy matching (case-insensitive, handles spaces/underscores)"""
+    if not col_name:
+        return None
+    
+    # Exact match first
+    if col_name in df.columns:
+        return col_name
+    
+    # Normalize the search term
+    search_term = col_name.lower().replace('_', '').replace(' ', '')
+    
+    # Try fuzzy matching
+    for col in df.columns:
+        normalized_col = str(col).lower().replace('_', '').replace(' ', '')
+        if search_term == normalized_col or search_term in normalized_col or normalized_col in search_term:
+            return col
+    
+    return None
+
 def create_plot(df, dfw, plan, describe_stats, selected_columns=None):
     plot_spec = plan.get("plot")
     if not plot_spec:
@@ -175,9 +195,12 @@ def create_plot(df, dfw, plan, describe_stats, selected_columns=None):
     ptype = plot_spec.get("type", "bar")
     title = plot_spec.get("title", "Chart")
     plot_df = df if describe_stats else dfw
-    x = plot_spec.get("x")
-    y = plot_spec.get("y")
     
+    # Use fuzzy matching to find columns
+    x = find_column(plot_df, plot_spec.get("x"))
+    y = find_column(plot_df, plot_spec.get("y"))
+    
+    # Fallback: auto-select columns if not found
     if not x and len(plot_df.columns) > 0:
         categorical_cols = plot_df.select_dtypes(include=['object', 'category']).columns
         x = categorical_cols[0] if len(categorical_cols) > 0 else plot_df.columns[0]
@@ -187,16 +210,16 @@ def create_plot(df, dfw, plan, describe_stats, selected_columns=None):
     
     try:
         if ptype == "pie" and x and x in plot_df.columns:
-            value_counts = plot_df[x].value_counts()
+            value_counts = plot_df[x].value_counts().head(10)  # Limit to top 10 for readability
             fig = go.Figure(data=[go.Pie(labels=value_counts.index, values=value_counts.values, hole=0.3)])
             fig.update_layout(title=title, width=900, height=500)
             return fig
         elif ptype == "bar" and x and x in plot_df.columns and y and y in plot_df.columns:
-            fig = px.bar(plot_df, x=x, y=y, title=title)
+            fig = px.bar(plot_df.head(50), x=x, y=y, title=title)  # Limit rows for performance
             fig.update_layout(width=900, height=500)
             return fig
         elif ptype == "line" and x and x in plot_df.columns and y and y in plot_df.columns:
-            fig = px.line(plot_df, x=x, y=y, title=title, markers=True)
+            fig = px.line(plot_df.head(100), x=x, y=y, title=title, markers=True)
             fig.update_layout(width=900, height=500)
             return fig
         elif ptype == "hist" and y and y in plot_df.columns:
@@ -204,10 +227,11 @@ def create_plot(df, dfw, plan, describe_stats, selected_columns=None):
             fig.update_layout(width=900, height=500)
             return fig
         elif ptype == "scatter" and x and x in plot_df.columns and y and y in plot_df.columns:
-            fig = px.scatter(plot_df, x=x, y=y, title=title)
+            fig = px.scatter(plot_df.head(200), x=x, y=y, title=title)  # Limit for performance
             fig.update_layout(width=900, height=500)
             return fig
-    except:
+    except Exception as e:
+        print(f"Error creating plot: {str(e)}")
         pass
     return None