Merge pull request #543 from aacedar/thread_safe
update python_execute thread safe
This commit is contained in:
commit
be5c2646af
@ -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)
|
||||
|
Loading…
x
Reference in New Issue
Block a user