Coder-Vippro/ChatGPT-Discord-Bot
1
0
Fork 0
Code Issues Pull Requests 3 Actions Packages Projects Releases 13 Wiki Activity

refactor: enhance Python code execution with async package installation and improved timeout handling

Browse Source
This commit is contained in:
cauvang32
2025-08-25 21:01:22 +07:00
parent ac6bb8c582
commit ecfc2b48d5
3 changed files with 83 additions and 62 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
src/utils/openai_utils.py
View File

@@ -184,19 +184,19 @@ def get_tools_for_model() -> List[Dict[str, Any]]:
"type": "function", "type": "function",
"function": { "function": {
"name": "execute_python_code", "name": "execute_python_code",
"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.", "description": "Execute Python code with automatic package installation. IMPORTANT: If your code imports any library (pandas, numpy, requests, matplotlib, etc.), you MUST include it in 'install_packages' parameter or the code will fail. Always use print() statements to show output. Examples of packages: numpy, pandas, matplotlib, seaborn, requests, beautifulsoup4, opencv-python, scikit-learn, plotly, etc.",
"parameters": { "parameters": {
"type": "object", "type": "object",
"properties": { "properties": {
"code": {"type": "string", "description": "Python code with print() statements"}, "code": {"type": "string", "description": "Python code with print() statements for output"},
"input_data": {"type": "string", "description": "Optional input data"}, "input_data": {"type": "string", "description": "Optional input data"},
"install_packages": { "install_packages": {
"type": "array", "type": "array",
"items": {"type": "string"}, "items": {"type": "string"},
"description": "List of pip package names to install before running code (e.g., ['requests', 'beautifulsoup4', 'opencv-python'])" "description": "REQUIRED: List ALL pip packages your code imports. Examples: ['pandas'] for pd.read_csv(), ['matplotlib'] for plt.plot(), ['requests'] for HTTP requests, ['numpy'] for arrays, ['beautifulsoup4'] for HTML parsing, etc. If you use ANY import statements, add the package here!"
}, },
"enable_visualization": {"type": "boolean", "description": "For charts/graphs"}, "enable_visualization": {"type": "boolean", "description": "For charts/graphs"},
"timeout": {"type": "integer", "default": 30, "minimum": 1, "maximum": 240} "timeout": {"type": "integer", "default": 60, "minimum": 1, "maximum": 300, "description": "Execution timeout in seconds (default 60, max 300)"}
}, },
"required": ["code"] "required": ["code"]
} }

137
src/utils/python_executor.py
View File

@@ -6,6 +6,7 @@ This module provides a completely secure isolated execution environment.
import os import os
import sys import sys
import subprocess import subprocess
import asyncio
import tempfile import tempfile
import venv import venv
import shutil import shutil
@@ -189,9 +190,9 @@ class SecureExecutor:
# For unknown packages, be restrictive # For unknown packages, be restrictive
return False, f"Package '{package}' is not in the approved safe list" return False, f"Package '{package}' is not in the approved safe list"
def install_packages_clean(self, packages: List[str], pip_path: str) -> Tuple[List[str], List[str]]: async def install_packages_clean(self, packages: List[str], pip_path: str) -> Tuple[List[str], List[str]]:
""" """
Install packages in the clean virtual environment. Install packages in the clean virtual environment (async to prevent blocking).
Args: Args:
packages: List of package names to install packages: List of package names to install
@@ -211,31 +212,40 @@ class SecureExecutor:
continue continue
try: try:
# Install package in the clean virtual environment # Install package in the clean virtual environment using async subprocess
result = subprocess.run( process = await asyncio.create_subprocess_exec(
[pip_path, "install", package], pip_path, "install", package,
capture_output=True, stdout=asyncio.subprocess.PIPE,
text=True, stderr=asyncio.subprocess.PIPE,
timeout=120, # 2 minutes per package cwd=self.temp_dir
check=False,
cwd=self.temp_dir # Run from temp directory
) )
if result.returncode == 0: try:
installed.append(package) stdout, stderr = await asyncio.wait_for(process.communicate(), timeout=120)
else: return_code = process.returncode
if return_code == 0:
installed.append(package)
else:
failed.append(package)
except asyncio.TimeoutError:
# Kill the process if it times out
try:
process.kill()
await process.wait()
except:
pass
failed.append(package) failed.append(package)
except subprocess.TimeoutExpired:
failed.append(package)
except Exception as e: except Exception as e:
failed.append(package) failed.append(package)
return installed, failed return installed, failed
def execute_code_secure(self, code: str, python_path: str, timeout: int) -> Dict[str, Any]: async def execute_code_secure(self, code: str, python_path: str, timeout: int) -> Dict[str, Any]:
""" """
Execute Python code in the completely isolated environment. Execute Python code in the completely isolated environment (async to prevent blocking).
Args: Args:
code: Python code to execute code: Python code to execute
@@ -254,14 +264,12 @@ class SecureExecutor:
with open(code_file, 'w', encoding='utf-8') as f: with open(code_file, 'w', encoding='utf-8') as f:
f.write(code) f.write(code)
# Execute code in completely isolated environment # Execute code in completely isolated environment using async subprocess
result = subprocess.run( process = await asyncio.create_subprocess_exec(
[python_path, code_file], python_path, code_file,
capture_output=True, stdout=asyncio.subprocess.PIPE,
text=True, stderr=asyncio.subprocess.PIPE,
timeout=timeout, cwd=self.temp_dir,
check=False,
cwd=self.temp_dir, # Run from isolated directory
env={ # Minimal environment variables env={ # Minimal environment variables
'PATH': os.path.dirname(python_path), 'PATH': os.path.dirname(python_path),
'PYTHONPATH': '', 'PYTHONPATH': '',
@@ -269,41 +277,54 @@ class SecureExecutor:
} }
) )
execution_time = time.time() - start_time try:
# Wait for process completion with timeout
# Process results stdout, stderr = await asyncio.wait_for(process.communicate(), timeout=timeout)
output = result.stdout return_code = process.returncode
error_output = result.stderr
execution_time = time.time() - start_time
# Truncate output if too large
if len(output) > MAX_OUTPUT_SIZE: # Process results
output = output[:MAX_OUTPUT_SIZE] + "\n... (output truncated)" output = stdout.decode('utf-8') if stdout else ""
error_output = stderr.decode('utf-8') if stderr else ""
if result.returncode == 0:
return { # Truncate output if too large
"success": True, if len(output) > MAX_OUTPUT_SIZE:
"output": output, output = output[:MAX_OUTPUT_SIZE] + "\n... (output truncated)"
"error": error_output if error_output else "",
"execution_time": execution_time, if return_code == 0:
"return_code": result.returncode return {
} "success": True,
else: "output": output,
"error": error_output if error_output else "",
"execution_time": execution_time,
"return_code": return_code
}
else:
return {
"success": False,
"output": output,
"error": error_output,
"execution_time": execution_time,
"return_code": return_code
}
except asyncio.TimeoutError:
# Kill the process if it times out
try:
process.kill()
await process.wait()
except:
pass
return { return {
"success": False, "success": False,
"output": output, "output": "",
"error": error_output, "error": f"Code execution timed out after {timeout} seconds",
"execution_time": execution_time, "execution_time": timeout,
"return_code": result.returncode "return_code": -1
} }
except subprocess.TimeoutExpired:
return {
"success": False,
"output": "",
"error": f"Code execution timed out after {timeout} seconds",
"execution_time": timeout,
"return_code": -1
}
except Exception as e: except Exception as e:
execution_time = time.time() - start_time execution_time = time.time() - start_time
error_msg = f"Execution error: {str(e)}" error_msg = f"Execution error: {str(e)}"
@@ -369,7 +390,7 @@ async def execute_python_code(args: Dict[str, Any]) -> Dict[str, Any]:
installed_packages = [] installed_packages = []
failed_packages = [] failed_packages = []
if packages_to_install: if packages_to_install:
installed_packages, failed_packages = executor.install_packages_clean(packages_to_install, pip_path) installed_packages, failed_packages = await executor.install_packages_clean(packages_to_install, pip_path)
# Prepare code with input data if provided # Prepare code with input data if provided
if input_data: if input_data:
@@ -379,7 +400,7 @@ async def execute_python_code(args: Dict[str, Any]) -> Dict[str, Any]:
code_with_input = code code_with_input = code
# Execute code in clean environment # Execute code in clean environment
result = executor.execute_code_secure(code_with_input, python_path, timeout) result = await executor.execute_code_secure(code_with_input, python_path, timeout)
# Add package installation info # Add package installation info
if installed_packages: if installed_packages:

0
src/utils/python_executor_new.py Normal file
View File