v113

app.py

@@ -170,8 +170,15 @@ def run_crewai_process(user_query, model, temperature):
         yield final_answer_chat, agent_thoughts, generated_code, execution_output, None, None
 
         # --- Execute the generated code ---
-        # 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
+        # Check for common plot filename patterns
+        # First check for symbol-specific plots (e.g., META_plot.png)
+        symbol_plot_patterns = ['META_plot.png', 'AAPL_plot.png', 'MSFT_plot.png', 'GOOG_plot.png', 'TSLA_plot.png']
+        
+        # Also check for generic plot filenames
+        generic_plot_patterns = ['plot.png', 'output.png', 'figure.png']
+        
+        # Combine all patterns to check
+        plot_file_paths = symbol_plot_patterns + generic_plot_patterns
 
         if generated_code:
             try:
@@ -215,22 +222,44 @@ def run_crewai_process(user_query, model, temperature):
 
                 # Check for the generated plot file in all possible locations
                 generated_plot_path = None
+                plot_found = False
+                
                 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
+                        
+                        # Add the plot to the execution output
+                        try:
+                            import base64
+                            with open(plot_abs_path, 'rb') as img_file:
+                                img_str = base64.b64encode(img_file.read()).decode('utf-8')
+                                execution_output += f"\n\n![Generated Plot](data:image/png;base64,{img_str})"
+                        except Exception as e:
+                            execution_output += f"\n\nNote: Could not embed plot in output: {str(e)}"
+                        
                         break
                 
-                if not generated_plot_path:
+                if not plot_found:
                     # 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')]
+                    png_files = [f for f in os.listdir(current_dir) if f.endswith('.png') and not f.startswith('gradio_')]
                     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}")
+                        
+                        # Add the plot to the execution output
+                        try:
+                            import base64
+                            with open(plot_abs_path, 'rb') as img_file:
+                                img_str = base64.b64encode(img_file.read()).decode('utf-8')
+                                execution_output += f"\n\n![Generated Plot](data:image/png;base64,{img_str})"
+                        except Exception as e:
+                            execution_output += f"\n\nNote: Could not embed plot in output: {str(e)}"
                     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"
@@ -250,14 +279,14 @@ def run_crewai_process(user_query, model, temperature):
         # Update final answer chat to reflect execution attempt
         execution_complete_msg = "Code execution finished. See 'Execution Output'."
         if generated_plot_path:
-            plot_msg = "Plot generated successfully. See the 'Generated Plot' tab below."
+            plot_msg = f"Plot generated successfully at: {generated_plot_path}"
             final_answer_chat = [
                 {"role": "user", "content": str(user_query)},
                 {"role": "assistant", "content": execution_complete_msg},
                 {"role": "assistant", "content": plot_msg}
             ]
+            # Return the plot path to be displayed
             yield final_answer_chat, agent_thoughts, generated_code, execution_output, None, generated_plot_path
-            return
         else:
             no_plot_msg = "No plot was generated. Make sure your query includes a request for a visualization. Check the 'Execution Output' tab for any errors."
             final_answer_chat = [
@@ -265,8 +294,8 @@ def run_crewai_process(user_query, model, temperature):
                 {"role": "assistant", "content": execution_complete_msg},
                 {"role": "assistant", "content": no_plot_msg}
             ]
+            # Return None for plot path
             yield final_answer_chat, agent_thoughts, generated_code, execution_output, None, None
-            return
 
         yield agent_thoughts, final_answer_chat, generated_code, execution_output, generated_plot_path