Update main.py

main.py

@@ -838,73 +838,92 @@ async def visualize_with_natural_language(
     style: str = Form("seaborn-v0_8")
 ):
     try:
-        # Read and process Excel file
+        # Read Excel file
         content = await file.read()
         df = pd.read_excel(BytesIO(content))
         
-        # Automatic type conversion
+        # Clean and convert data types
         for col in df.columns:
-            # Try converting to numeric first
+            # First try numeric conversion
             df[col] = pd.to_numeric(df[col], errors='ignore')
-            # Then try dates
+            # Then try datetime conversion
             try:
                 df[col] = pd.to_datetime(df[col], errors='ignore')
             except:
                 pass
+            # Finally clean any remaining strings
+            df[col] = df[col].astype(str).str.strip().replace('nan', np.nan)
         
         # Generate visualization request
         vis_request = interpret_natural_language(prompt, df.columns.tolist())
+        if not vis_request:
+            raise HTTPException(400, "Could not interpret visualization request")
         
-        # Create plot with type-safe data
+        # Create visualization with type safety
         plt.style.use(style)
         fig, ax = plt.subplots(figsize=(10, 6))
         
-        if vis_request.chart_type == "heatmap":
-            numeric_df = df.select_dtypes(include=['number'])
-            sns.heatmap(numeric_df.corr(), annot=True)
-        else:
-            # Type-safe plotting
-            plot_data = df.copy()
-            if vis_request.x_column:
-                plot_data[vis_request.x_column] = pd.to_numeric(
-                    plot_data[vis_request.x_column], 
-                    errors='ignore'
-                )
-            if vis_request.y_column:
-                plot_data[vis_request.y_column] = pd.to_numeric(
-                    plot_data[vis_request.y_column],
-                    errors='ignore'
-                )
-            
-            if vis_request.chart_type == "line":
-                sns.lineplot(
-                    data=plot_data,
-                    x=vis_request.x_column,
-                    y=vis_request.y_column,
-                    hue=vis_request.hue_column
-                )
-            elif vis_request.chart_type == "bar":
-                sns.barplot(
-                    data=plot_data,
-                    x=vis_request.x_column,
-                    y=vis_request.y_column,
-                    hue=vis_request.hue_column
-                )
-            # Add other chart types similarly...
-        
-        plt.title(vis_request.title)
-        buffer = BytesIO()
-        plt.savefig(buffer, format='png')
-        plt.close()
-        
-        return {
-            "image": base64.b64encode(buffer.getvalue()).decode('utf-8'),
-            "chart_type": vis_request.chart_type,
-            "columns": list(df.columns)
-        }
+        try:
+            if vis_request.chart_type == "heatmap":
+                numeric_df = df.select_dtypes(include=['number'])
+                if numeric_df.empty:
+                    raise ValueError("No numeric columns for heatmap")
+                sns.heatmap(numeric_df.corr(), annot=True)
+            else:
+                # Ensure numeric data for plotting
+                plot_data = df.copy()
+                if vis_request.x_column:
+                    plot_data[vis_request.x_column] = pd.to_numeric(
+                        plot_data[vis_request.x_column],
+                        errors='coerce'
+                    )
+                if vis_request.y_column:
+                    plot_data[vis_request.y_column] = pd.to_numeric(
+                        plot_data[vis_request.y_column],
+                        errors='coerce'
+                    )
+                
+                # Remove rows with missing numeric data
+                plot_data = plot_data.dropna()
+                
+                if vis_request.chart_type == "line":
+                    sns.lineplot(
+                        data=plot_data,
+                        x=vis_request.x_column,
+                        y=vis_request.y_column,
+                        hue=vis_request.hue_column
+                    )
+                elif vis_request.chart_type == "bar":
+                    sns.barplot(
+                        data=plot_data,
+                        x=vis_request.x_column,
+                        y=vis_request.y_column,
+                        hue=vis_request.hue_column
+                    )
+                # Add other chart types as needed...
         
+            plt.title(vis_request.title)
+            buffer = BytesIO()
+            plt.savefig(buffer, format='png', bbox_inches='tight')
+            plt.close()
+            
+            return {
+                "status": "success",
+                "image": base64.b64encode(buffer.getvalue()).decode('utf-8'),
+                "chart_type": vis_request.chart_type,
+                "columns": list(df.columns),
+                "x_column": vis_request.x_column,
+                "y_column": vis_request.y_column,
+                "hue_column": vis_request.hue_column
+            }
+            
+        except Exception as e:
+            raise HTTPException(400, f"Plotting error: {str(e)}")
+
+    except HTTPException:
+        raise
     except Exception as e:
-        raise HTTPException(400, f"Visualization failed: {str(e)}")
+        raise HTTPException(500, f"Server error: {str(e)}")
 
 
 @app.get("/visualize/styles")