Update main.py

main.py

@@ -906,7 +906,14 @@ async def natural_language_visualization(
                         ax=ax
                     )
                 # Add other chart types as needed...
-
+   buf = BytesIO()
+        plt.savefig(buf, format='png', bbox_inches='tight')
+        plt.close(fig)
+        buf.seek(0)
+        
+        # Encode without newlines
+        image_base64 = base64.b64encode(buf.read()).decode('utf-8')
+        
             ax.set_title(vis_request.title)
             buf = BytesIO()
             plt.savefig(buf, format='png', bbox_inches='tight')
@@ -921,6 +928,7 @@ async def natural_language_visualization(
                 "x_column": vis_request.x_column,
                 "y_column": vis_request.y_column,
                 "hue_column": vis_request.hue_column
+                "code": generated_code
             }
             
         except Exception as e: