v116

app.py

@@ -200,6 +200,9 @@ def run_crewai_process(user_query, model, temperature):
                 
                 # Add helper functions at the beginning of the script
                 helpers = """
+# Standard plot filename to use
+PLOT_FILENAME = 'generated_plot.png'
+
 # Helper functions for data processing
 def safe_get_column(df, column):
     # Try exact match first
@@ -221,13 +224,22 @@ def safe_get_column(df, column):
     # If still not found, raise a helpful error
     raise KeyError(f"Column '{column}' not found in DataFrame. Available columns: {list(df.columns)}")
 
-def plot_and_save(plt, filename='plot.png'):
+def plot_and_save(plt, filename=None):
     import os
+    if filename is None:
+        filename = PLOT_FILENAME
+    
+    # Convert to absolute path
+    filename = os.path.abspath(filename)
+    
     # Ensure the directory exists
-    os.makedirs(os.path.dirname(os.path.abspath(filename)) or '.', exist_ok=True)
-    plt.savefig(filename)
+    os.makedirs(os.path.dirname(filename) or '.', exist_ok=True)
+    
+    # Save the plot
+    plt.savefig(filename, bbox_inches='tight', dpi=100)
     plt.close()
-    print(f"Plot saved as {os.path.abspath(filename)}")
+    print(f"[DEBUG] Plot saved to: {filename}")
+    return filename
 
 # Monkey patch DataFrame to add safe column access
 import pandas as pd
@@ -255,6 +267,10 @@ pd.DataFrame.safe_get = safe_get_column
                     script_content
                 )
                 
+                # If no savefig call is found, add one at the end of the script
+                if 'plt.savefig(' not in script_content and 'plot_and_save(' not in script_content:
+                    script_content += "\n# Ensure plot is saved\nif 'plt' in locals() and len(plt.get_fignums()) > 0:\n    plot_and_save(plt, PLOT_FILENAME)\n"
+                
                 # Write the updated script back
                 with open(temp_script_path, 'w') as f:
                     f.write(script_content)
@@ -281,12 +297,31 @@ pd.DataFrame.safe_get = safe_get_column
                 generated_plot_path = None
                 plot_found = False
                 
+                # First check the standard plot filename
+                plot_file_paths = ['generated_plot.png', 'plot.png', 'META_plot.png', 'AAPL_plot.png', 'MSFT_plot.png', 'output.png']
+                
+                # Also check for any .png files in the current directory
+                current_dir = os.path.abspath('.')
+                png_files = [f for f in os.listdir(current_dir) 
+                           if f.endswith('.png') and not f.startswith('gradio_')]
+                
+                # Add any found .png files to our search paths
+                plot_file_paths.extend(png_files)
+                
+                # Make paths absolute and remove duplicates
+                plot_file_paths = list(dict.fromkeys([os.path.abspath(f) for f in plot_file_paths]))
+                
+                print(f"[DEBUG] Looking for plot files in: {plot_file_paths}")
+                
                 for plot_file in plot_file_paths:
-                    plot_abs_path = os.path.abspath(plot_file)
-                    if os.path.exists(plot_abs_path):
-                        print(f"Plot file found at: {plot_abs_path}")
-                        generated_plot_path = plot_abs_path
-                        plot_found = True
+                    try:
+                        if os.path.exists(plot_file) and os.path.getsize(plot_file) > 0:
+                            print(f"[DEBUG] Found plot file: {plot_file}")
+                            generated_plot_path = plot_file
+                            plot_found = True
+                            break
+                    except Exception as e:
+                        print(f"[DEBUG] Error checking plot file {plot_file}: {e}")
                         
                         # Add the plot to the execution output
                         try: