refactor: enhance code execution and analysis features with package installation support and detailed output display

README.md

@@ -9,6 +9,9 @@ Welcome to **ChatGPT Discord Bot**! This bot provides a powerful AI assistant fo
 - **Image Generation**: Creates custom images from text prompts using Runware's API
 - **Data Analysis**: Analyzes CSV and Excel files with visualizations (distributions, correlations, box plots, etc.)
 - **Code Interpretation**: Executes Python code for calculations and data processing
+  - **Package Installation**: Automatically installs required Python packages for code execution
+  - **Code Display**: Shows executed code, input, and output in Discord chat for transparency
+  - **Secure Sandbox**: Runs code in a controlled environment with safety restrictions
 - **Reminder System**: Sets timed reminders with custom timezone support
 - **Web Tools**:
   - **Google Search**: Searches the web and provides relevant information

src/module/message_handler.py

@@ -214,6 +214,9 @@ class MessageHandler:
             if not user_id:
                 user_id = self._find_user_id_from_current_task()
             
+            # Get the Discord message to send code execution display
+            discord_message = self._get_discord_message_from_current_task()
+            
             # Add file context if user has uploaded data files
             if user_id and user_id in self.user_data_files:
                 file_info = self.user_data_files[user_id]
@@ -227,14 +230,80 @@ class MessageHandler:
                 
                 logging.info(f"Added file context to Python execution for user {user_id}")
             
+            # Extract code, input, and packages for display
+            code_to_execute = args.get("code", "")
+            input_data = args.get("input_data", "")
+            packages_to_install = args.get("install_packages", [])
+            
             # Import and call Python executor
             from src.utils.python_executor import execute_python_code
             execute_result = await execute_python_code(args)
             
+            # Display the executed code information in Discord (but not save to history)
+            if discord_message and code_to_execute:
+                try:
+                    # Create a formatted code execution display as a separate message
+                    execution_display = "**🐍 Python Code Execution**\n\n"
+                    
+                    # Show packages to install if any
+                    if packages_to_install:
+                        execution_display += f"**📦 Installing packages:** {', '.join(packages_to_install)}\n\n"
+                    
+                    # Show input data if any
+                    if input_data:
+                        execution_display += "**📥 Input:**\n```\n"
+                        execution_display += input_data[:500]  # Limit input length
+                        if len(input_data) > 500:
+                            execution_display += "\n... (input truncated)"
+                        execution_display += "\n```\n\n"
+                    
+                    # Show the actual code
+                    execution_display += "**💻 Code:**\n```python\n"
+                    # Clean up the code display (remove file context comments)
+                    code_lines = code_to_execute.split('\n')
+                    clean_code_lines = []
+                    for line in code_lines:
+                        if not (line.strip().startswith('# Data file available:') or 
+                               line.strip().startswith('# File path:') or 
+                               line.strip().startswith('# You can access this file using:')):
+                            clean_code_lines.append(line)
+                    
+                    clean_code = '\n'.join(clean_code_lines).strip()
+                    execution_display += clean_code[:1000]  # Limit code length for Discord
+                    if len(clean_code) > 1000:
+                        execution_display += "\n... (code truncated)"
+                    execution_display += "\n```\n\n"
+                    
+                    # Show the output
+                    if execute_result and execute_result.get("success"):
+                        output = execute_result.get("output", "")
+                        # Remove package installation info from output if it exists
+                        if output and "Installed packages:" in output:
+                            lines = output.split('\n')
+                            output = '\n'.join(lines[2:]) if len(lines) > 2 else ""
+                        
+                        if output and output.strip():
+                            execution_display += "**📤 Output:**\n```\n"
+                            execution_display += output[:1000]  # Limit output length for Discord
+                            if len(output) > 1000:
+                                execution_display += "\n... (output truncated)"
+                            execution_display += "\n```"
+                        else:
+                            execution_display += "**📤 Output:** *(No output)*"
+                    else:
+                        error_msg = execute_result.get("error", "Unknown error") if execute_result else "Execution failed"
+                        execution_display += f"**❌ Error:**\n```\n{error_msg[:800]}\n```"
+                        if len(error_msg) > 800:
+                            execution_display += "*(Error message truncated)*"
+                    
+                    # Send the execution display to Discord as a separate message
+                    await discord_message.channel.send(execution_display)
+                    
+                except Exception as e:
+                    logging.error(f"Error displaying code execution: {str(e)}")
+            
             # If there are visualizations, handle them
             if execute_result and execute_result.get("visualizations"):
-                discord_message = self._get_discord_message_from_current_task()
-                
                 for i, viz_path in enumerate(execute_result["visualizations"]):
                     try:
                         with open(viz_path, 'rb') as f:
@@ -314,14 +383,66 @@ class MessageHandler:
             # Add user_id to args for the data analyzer
             args["user_id"] = user_id
             
+            # Get the Discord message to send code execution display
+            discord_message = self._get_discord_message_from_current_task()
+            
             # Import and call data analyzer
             from src.utils.data_analyzer import analyze_data_file
             result = await analyze_data_file(args)
             
+            # Display the generated code if available
+            if discord_message and result and result.get("generated_code"):
+                try:
+                    # Create a formatted code execution display as a separate message
+                    execution_display = "**<2A> Data Analysis Execution**\n\n"
+                    
+                    # Show the file being analyzed
+                    file_path = args.get("file_path", "")
+                    if file_path:
+                        filename = os.path.basename(file_path)
+                        execution_display += f"**📁 Analyzing file:** `{filename}`\n\n"
+                    
+                    # Show the analysis type if specified
+                    analysis_type = args.get("analysis_type", "")
+                    custom_analysis = args.get("custom_analysis", "")
+                    if analysis_type:
+                        execution_display += f"**🔍 Analysis type:** {analysis_type}\n\n"
+                    if custom_analysis:
+                        execution_display += f"**📝 Custom request:** {custom_analysis}\n\n"
+                    
+                    # Show the generated code
+                    execution_display += "**💻 Generated Code:**\n```python\n"
+                    generated_code = result["generated_code"]
+                    execution_display += generated_code[:1000]  # Limit code length for Discord
+                    if len(generated_code) > 1000:
+                        execution_display += "\n... (code truncated)"
+                    execution_display += "\n```\n\n"
+                    
+                    # Show the output
+                    if result.get("success"):
+                        output = result.get("output", "")
+                        if output and output.strip():
+                            execution_display += "**📊 Analysis Results:**\n```\n"
+                            execution_display += output[:1000]  # Limit output length for Discord
+                            if len(output) > 1000:
+                                execution_display += "\n... (output truncated)"
+                            execution_display += "\n```"
+                        else:
+                            execution_display += "**📊 Analysis Results:** *(No text output - check for visualizations below)*"
+                    else:
+                        error_msg = result.get("error", "Unknown error")
+                        execution_display += f"**❌ Error:**\n```\n{error_msg[:800]}\n```"
+                        if len(error_msg) > 800:
+                            execution_display += "*(Error message truncated)*"
+                    
+                    # Send the execution display to Discord as a separate message
+                    await discord_message.channel.send(execution_display)
+                    
+                except Exception as e:
+                    logging.error(f"Error displaying data analysis code: {str(e)}")
+            
             # If there are visualizations, handle them for Discord
             if result and result.get("visualizations"):
-                discord_message = self._get_discord_message_from_current_task()
-                
                 for i, viz_path in enumerate(result["visualizations"]):
                     try:
                         with open(viz_path, 'rb') as f:

src/utils/openai_utils.py

@@ -184,13 +184,17 @@ def get_tools_for_model() -> List[Dict[str, Any]]:
             "type": "function",            
             "function": {
                 "name": "execute_python_code",
-                "description": "Execute Python code. MUST use print() for output.",
+                "description": "Execute Python code with package installation support. If you need specific packages, list them in 'install_packages' parameter. The system will install them before running your code. MUST use print() for output.",
                 "parameters": {
                     "type": "object",
                     "properties": {
                         "code": {"type": "string", "description": "Python code with print() statements"},
                         "input_data": {"type": "string", "description": "Optional input data"},
-                        "install_packages": {"type": "array", "items": {"type": "string"}},
+                        "install_packages": {
+                            "type": "array", 
+                            "items": {"type": "string"}, 
+                            "description": "List of pip package names to install before running code (e.g., ['requests', 'beautifulsoup4', 'opencv-python'])"
+                        },
                         "enable_visualization": {"type": "boolean", "description": "For charts/graphs"},
                         "timeout": {"type": "integer", "default": 30, "minimum": 1, "maximum": 240}
                     },

src/utils/python_executor.py

@@ -156,7 +156,7 @@ async def execute_python_code(args: Dict[str, Any]) -> Dict[str, Any]:
         args: Dictionary containing:
             - code: The Python code to execute
             - input: Optional input data for the code
-            - install_packages: Optional list of packages to install
+            - install_packages: List of packages to install before execution
             - timeout: Optional timeout in seconds (default: 30)
             
     Returns:
@@ -175,11 +175,20 @@ async def execute_python_code(args: Dict[str, Any]) -> Dict[str, Any]:
                 "output": ""
             }
         
-        # Install packages if requested
+        # Install requested packages first
+        installed_packages = []
         if packages_to_install:
+            logger.info(f"Installing requested packages: {packages_to_install}")
             install_result = await install_packages(packages_to_install)
-            if not install_result["success"]:
+            
-                logger.warning(f"Package installation issues: {install_result}")
+            if install_result["installed"]:
+                installed_packages = install_result["installed"]
+                logger.info(f"Successfully installed: {installed_packages}")
+            
+            if install_result["failed"]:
+                failed_packages = [f["package"] for f in install_result["failed"]]
+                logger.warning(f"Failed to install: {failed_packages}")
+                # Continue execution even if some packages failed to install
         
         # Sanitize the code
         is_safe, sanitized_code = sanitize_python_code(code)
@@ -197,6 +206,14 @@ async def execute_python_code(args: Dict[str, Any]) -> Dict[str, Any]:
         # Execute code in controlled environment
         result = await execute_code_safely(sanitized_code, input_data, timeout)
         
+        # Add information about installed packages to the result
+        if installed_packages:
+            result["installed_packages"] = installed_packages
+            # Prepend package installation info to output
+            if result.get("success"):
+                package_info = f"[Installed packages: {', '.join(installed_packages)}]\n\n"
+                result["output"] = package_info + result.get("output", "")
+        
         return result
         
     except Exception as e:
@@ -337,6 +354,22 @@ async def execute_code_safely(code: str, input_data: str, timeout: int) -> Dict[
                 "output": stdout_capture.getvalue(),
                 "stderr": stderr_capture.getvalue()
             }
+        except Exception as e:
+            # Capture execution errors with helpful information
+            error_msg = str(e)
+            stderr_content = stderr_capture.getvalue()
+            
+            # If it's an import error, provide helpful guidance
+            if "ModuleNotFoundError" in error_msg or "ImportError" in error_msg:
+                error_msg += "\n\nHint: If you need additional packages, specify them in the 'install_packages' parameter."
+            
+            return {
+                "success": False,
+                "error": f"Execution error: {error_msg}",
+                "output": stdout_capture.getvalue(),
+                "stderr": stderr_content + f"\nExecution error: {error_msg}",
+                "traceback": traceback.format_exc()
+            }
         
         # Restore stdout and stderr
         sys.stdout = old_stdout