Merge pull request #543 from aacedar/thread_safe

update python_execute thread safe

app/tool/python_execute.py

@@ -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():
+        with multiprocessing.Manager() as manager:
-            try:
+            result = manager.dict({
-                safe_globals = {"__builtins__": dict(__builtins__)}
+                "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
+            # timeout process
-                from io import StringIO
+            if proc.is_alive():
-
+                proc.terminate()
-                output_buffer = StringIO()
+                proc.join(1)
-                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():
                 return {
                     "observation": f"Execution timeout after {timeout} seconds",
                     "success": False,
                 }
-
+            return dict(result)
-        return result