Merge pull request #543 from aacedar/thread_safe

update python_execute thread safe
This commit is contained in:
Isaac 2025-03-13 21:15:17 +08:00 committed by GitHub
commit be5c2646af
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -1,4 +1,6 @@
import threading
import sys
from io import StringIO
import multiprocessing
from typing import Dict
from app.tool.base import BaseTool
@ -20,6 +22,20 @@ class PythonExecute(BaseTool):
"required": ["code"],
}
def _run_code(self, code: str, result_dict: dict, safe_globals: dict) -> None:
original_stdout = sys.stdout
try:
output_buffer = StringIO()
sys.stdout = output_buffer
exec(code, safe_globals, safe_globals)
result_dict["observation"] = output_buffer.getvalue()
result_dict["success"] = True
except Exception as e:
result_dict["observation"] = str(e)
result_dict["success"] = False
finally:
sys.stdout = original_stdout
async def execute(
self,
code: str,
@ -35,36 +51,29 @@ class PythonExecute(BaseTool):
Returns:
Dict: Contains 'output' with execution output or error message and 'success' status.
"""
result = {"observation": ""}
def run_code():
try:
safe_globals = {"__builtins__": dict(__builtins__)}
with multiprocessing.Manager() as manager:
result = manager.dict({
"observation": "",
"success": False
})
if isinstance(__builtins__, dict):
safe_globals = {"__builtins__": __builtins__}
else:
safe_globals = {"__builtins__": __builtins__.__dict__.copy()}
proc = multiprocessing.Process(
target=self._run_code,
args=(code, result, safe_globals)
)
proc.start()
proc.join(timeout)
import sys
from io import StringIO
output_buffer = StringIO()
sys.stdout = output_buffer
exec(code, safe_globals, {})
sys.stdout = sys.__stdout__
result["observation"] = output_buffer.getvalue()
except Exception as e:
result["observation"] = str(e)
result["success"] = False
thread = threading.Thread(target=run_code)
thread.start()
thread.join(timeout)
if thread.is_alive():
# timeout process
if proc.is_alive():
proc.terminate()
proc.join(1)
return {
"observation": f"Execution timeout after {timeout} seconds",
"success": False,
}
return result
return dict(result)