diff --git a/app/tool/python_execute.py b/app/tool/python_execute.py index 08ceffa85..110c0cd40 100644 --- a/app/tool/python_execute.py +++ b/app/tool/python_execute.py @@ -1,3 +1,4 @@ +import asyncio import multiprocessing import sys from io import StringIO @@ -36,22 +37,9 @@ def _run_code(self, code: str, result_dict: dict, safe_globals: dict) -> None: finally: sys.stdout = original_stdout - async def execute( - self, - code: str, - timeout: int = 5, - ) -> Dict: - """ - Executes the provided Python code with a timeout. - - Args: - code (str): The Python code to execute. - timeout (int): Execution timeout in seconds. - - Returns: - Dict: Contains 'output' with execution output or error message and 'success' status. - """ - + def _run_in_process(self, code: str, timeout: int) -> Dict: + """Run code in a subprocess. This method is blocking and should be + called via run_in_executor to avoid blocking the event loop.""" with multiprocessing.Manager() as manager: result = manager.dict({"observation": "", "success": False}) if isinstance(__builtins__, dict): @@ -64,7 +52,6 @@ async def execute( proc.start() proc.join(timeout) - # timeout process if proc.is_alive(): proc.terminate() proc.join(1) @@ -73,3 +60,21 @@ async def execute( "success": False, } return dict(result) + + async def execute( + self, + code: str, + timeout: int = 5, + ) -> Dict: + """ + Executes the provided Python code with a timeout. + + Args: + code (str): The Python code to execute. + timeout (int): Execution timeout in seconds. + + Returns: + Dict: Contains 'observation' with execution output or error message and 'success' status. + """ + loop = asyncio.get_event_loop() + return await loop.run_in_executor(None, self._run_in_process, code, timeout)