added message

codeexecutor.py

@@ -1,8 +1,8 @@
 import os
 import re
-import signal
 import subprocess
 import tempfile
+import multiprocessing
 from collections import Counter
 from contextlib import contextmanager
 from dataclasses import dataclass
@@ -12,17 +12,37 @@ class PythonREPL:
     def __init__(self, timeout=5):
         self.timeout = timeout
 
-    @contextmanager
-    def time_limit(self, seconds):
-        def signal_handler(*_):
-            raise TimeoutError(f"Timed out after {seconds} seconds.")
-
-        signal.signal(signal.SIGALRM, signal_handler)
-        signal.alarm(seconds)
-        try:
-            yield
-        finally:
-            signal.alarm(0)
+    @staticmethod
+    def _run_code(temp_file_path):
+        result = subprocess.run(
+            ["python3", temp_file_path],
+            capture_output=True,
+            check=False,
+            text=True
+        )
+        if result.returncode == 0:
+            return True, result.stdout.strip()
+        else:
+            error_msg = result.stderr.strip()
+            msgs = error_msg.split("\n")
+            new_msgs = []
+            want_next = False
+            for m in msgs:
+                if "Traceback" in m:
+                    new_msgs.append(m)
+                elif m == msgs[-1]:
+                    new_msgs.append(m)
+                elif temp_file_path in m:
+                    st = m.index('"/') + 1 if '"/' in m else 0
+                    ed = m.index(temp_file_path) + 1 if temp_file_path in m else None
+                    clr = m[st:ed] if not ed else m[st:]
+                    m = m.replace(clr, "")
+                    new_msgs.append(m)
+                    want_next = True
+                elif want_next:
+                    new_msgs.append(m)
+                    want_next = False
+            return False, "\n".join(new_msgs).strip()
 
     def __call__(self, query):
         query = "import math\nimport numpy as np\nimport sympy as sp\n" + query
@@ -32,43 +52,21 @@ class PythonREPL:
                 query[-1] = query[-1].split("#")[0]
             query[-1] = "print(" + query[-1] + ")"
         query = "\n".join(query)
+        
         with tempfile.TemporaryDirectory() as temp_dir:
             temp_file_path = os.path.join(temp_dir, "tmp.py")
             with open(temp_file_path, "w", encoding="utf-8") as f:
                 f.write(query)
-            with self.time_limit(self.timeout):
-                result = subprocess.run(
-                    ["python3", temp_file_path],
-                    capture_output=True,
-                    check=False,
-                    text=True,
-                    timeout=self.timeout,
-                )
-                if result.returncode == 0:
-                    output = result.stdout
-                    return True, output.strip()
-                error_msg = result.stderr.strip()
-                msgs = error_msg.split("\n")
-                new_msgs = []
-                want_next = False
-                for m in msgs:
-                    if "Traceback" in m:
-                        new_msgs.append(m)
-                    elif m == msgs[-1]:
-                        new_msgs.append(m)
-                    elif temp_file_path in m:
-                        st = m.index('"/') + 1 if '"/' in m else 0
-                        ed = m.index(temp_file_path) + 1 if temp_file_path in m else None
-                        clr = m[st:ed] if not ed else m[st:]
-                        m = m.replace(clr, "")
-                        new_msgs.append(m)
-                        want_next = True
-                    elif want_next:
-                        new_msgs.append(m)
-                        want_next = False
-                error_msg = "\n".join(new_msgs)
-                return False, error_msg.strip()
-            
+
+            with multiprocessing.Pool(1) as pool:
+                result = pool.apply_async(self._run_code, (temp_file_path,))
+                try:
+                    success, output = result.get(self.timeout)
+                except multiprocessing.TimeoutError:
+                    pool.terminate()
+                    return False, f"Timed out after {self.timeout} seconds."
+        return success, output
+
 
 def execute_completion(executor, completion, return_status, last_code_block):
     executions = re.findall(r"```python(.*?)```", completion, re.DOTALL)
@@ -108,9 +106,10 @@ def postprocess_completion(text, return_status, last_code_block):
     del executor
     return result
 
+
 def get_majority_vote(answers):
     if not len(answers):
         return 0
     c = Counter(answers)
     value, _ = c.most_common()[0]
-    return value
+    return value