refactor: enhance code execution and analysis features with package installation support and detailed output display
This commit is contained in:
@@ -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
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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}
|
||||
},
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user