v112

app.py

@@ -170,7 +170,8 @@ def run_crewai_process(user_query, model, temperature):
         yield final_answer_chat, agent_thoughts, generated_code, execution_output, None, None
 
         # --- Execute the generated code ---
-        plot_file_path = 'plot.png'  # Expected plot file name
+        # Check for both plot.png and {symbol}_plot.png
+        plot_file_paths = ['plot.png', 'META_plot.png', 'AAPL_plot.png', 'MSFT_plot.png']  # Common plot filenames
 
         if generated_code:
             try:
@@ -212,15 +213,27 @@ def run_crewai_process(user_query, model, temperature):
                 elif "FileNotFoundError: [Errno 2] No such file or directory: 'plot.png'" in execution_output:
                     execution_output += "\n\nPlot Generation Issue: The script ran, but the plot file was not created. Ensure the generated code includes commands to save the plot to 'plot.png'."
 
-                # Check for the generated plot file
-                plot_abs_path = os.path.abspath(plot_file_path)
-                if os.path.exists(plot_abs_path):
-                    print(f"Plot file found at: {plot_abs_path}")
-                    # Return the absolute path to ensure Gradio can find the file
-                    generated_plot_path = plot_abs_path
-                else:
-                    print(f"Plot file not found at expected path: {plot_abs_path}")
-                    execution_output += f"\nPlot file '{plot_abs_path}' not found after execution.\n\nMake sure the generated code includes:\n1. `plt.savefig('plot.png')` to save the plot\n2. `plt.close()` to close the figure after saving"
+                # Check for the generated plot file in all possible locations
+                generated_plot_path = None
+                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
+                        break
+                
+                if not generated_plot_path:
+                    # If no plot file was found, check the current directory for any .png files
+                    current_dir = os.path.abspath('.')
+                    png_files = [f for f in os.listdir(current_dir) if f.endswith('.png')]
+                    if png_files:
+                        # Use the first .png file found
+                        plot_abs_path = os.path.abspath(png_files[0])
+                        generated_plot_path = plot_abs_path
+                        print(f"Using plot file found at: {plot_abs_path}")
+                    else:
+                        print(f"No plot file found in {current_dir}")
+                        execution_output += "\nNo plot file was found after execution.\n\nMake sure the generated code includes:\n1. `plt.savefig('plot.png')` to save the plot\n2. `plt.close()` to close the figure after saving"
 
             except Exception as e:
                 traceback_str = traceback.format_exc()
@@ -357,15 +370,24 @@ def create_interface():
             agent_thoughts,
             generated_code,
             execution_output,
-            plot_output,
-            image_output
+            None,  # This is for plot_output (we'll handle it separately)
+            None   # This is for image_output (we'll handle it separately)
         ]
         
+        def process_results(chat, thoughts, code, output, plot_path):
+            # This function will be called after run_crewai_process
+            # Show the image in the image_output component
+            return chat, thoughts, code, output, gr.update(visible=plot_path is not None), gr.update(value=plot_path, visible=plot_path is not None)
+        
         submit_btn.click(
             fn=run_crewai_process,
             inputs=inputs,
             outputs=outputs,
             api_name="analyze"
+        ).then(
+            fn=process_results,
+            inputs=[final_answer_chat, agent_thoughts, generated_code, execution_output, image_output],
+            outputs=[final_answer_chat, agent_thoughts, generated_code, execution_output, plot_output, image_output]
         )
     
     return interface