Update main.py

main.py

@@ -410,26 +410,27 @@ async def visualize_with_code(
     filters: Optional[str] = Form(None)
 ):
     try:
-        # Validate file
         file_ext, content = await process_uploaded_file(file)
+
         if file_ext not in {"xlsx", "xls"}:
-            raise HTTPException(400, "Only Excel files are supported for visualization")
-        
-        # Read Excel file
+            raise HTTPException(400, "Visualization is only supported for Excel files")
+
         df = pd.read_excel(io.BytesIO(content))
-        
-        # Parse filters if provided
-        filter_dict = {}
+
+        if df.empty:
+            raise HTTPException(400, "The uploaded Excel file is empty")
+
+        # Convert filters from string to dictionary safely
+        filters_dict = None
         if filters:
             try:
-                filter_dict = ast.literal_eval(filters)
-                if not isinstance(filter_dict, dict):
-                    filter_dict = {}
-            except:
-                filter_dict = {}
-        
-        # Create visualization request
-        vis_request = VisualizationRequest(
+                filters_dict = ast.literal_eval(filters)
+                if not isinstance(filters_dict, dict):
+                    raise ValueError()
+            except Exception:
+                raise HTTPException(400, "Invalid format for filters. Must be a valid dictionary string.")
+
+        viz_request = VisualizationRequest(
             chart_type=chart_type,
             x_column=x_column,
             y_column=y_column,
@@ -438,41 +439,17 @@ async def visualize_with_code(
             x_label=x_label,
             y_label=y_label,
             style=style,
-            filters=filter_dict
+            filters=filters_dict
         )
-        
-        # Generate visualization code
-        visualization_code = generate_visualization_code(df, vis_request)
-        
-        # Execute the code to generate the plot
-        plt.figure()
-        local_vars = {}
-        exec(visualization_code, globals(), local_vars)
-        
-        # Save the plot to a temporary file
-        with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmpfile:
-            plt.savefig(tmpfile.name, format='png', dpi=300)
-            plt.close()
-            
-            # Read the image back as bytes
-            with open(tmpfile.name, "rb") as f:
-                image_bytes = f.read()
-        
-        # Encode image as base64
-        image_base64 = base64.b64encode(image_bytes).decode('utf-8')
-        
-        return {
-            "status": "success",
-            "image": f"data:image/png;base64,{image_base64}",
-            "code": visualization_code,
-            "data_preview": df.head().to_dict(orient='records')
-        }
-        
+
+        code = generate_visualization_code(df, viz_request)
+        return {"code": code}
+
     except HTTPException:
         raise
     except Exception as e:
-        logger.error(f"Visualization failed: {str(e)}\n{traceback.format_exc()}")
-        raise HTTPException(500, detail=f"Visualization failed: {str(e)}")
+        logger.error(f"Visualization code generation failed: {str(e)}")
+        raise HTTPException(500, f"Visualization code generation failed: {str(e)}")
 
 @app.post("/visualize/natural")
 @limiter.limit("5/minute")