Update tools/python_runner.py

tools/python_runner.py

@@ -1,57 +1,53 @@
 from typing import Any, Dict
+from smolagents.tools import Tool
 import io
-import traceback
 import contextlib
+import traceback
 import signal
 
-# ---- your safe executor ----
-
 class TimeoutException(Exception):
     pass
 
 def timeout_handler(signum, frame):
     raise TimeoutException("Execution timed out.")
 
-def run_python(code: str, timeout: int = 5) -> Dict[str, Any]:
-    safe_globals = {
-        "__builtins__": {
-            "print": print,
-            "range": range,
-            "len": len,
-            "int": int,
-            "float": float,
-            "str": str,
-            "list": list,
-            "dict": dict,
-            "sum": sum,
-            "min": min,
-            "max": max,
-            "abs": abs,
+class PythonRunnerTool(Tool):
+    name = "python_runner"
+    description = "Executes Python code safely and returns stdout or errors."
+    inputs = {'code': {'type': 'string', 'description': 'Python code to execute'}}
+    output_type = "string"
+
+    def forward(self, code: str) -> str:
+        stdout_buffer = io.StringIO()
+        safe_globals = {
+            "__builtins__": {
+                "print": print,
+                "range": range,
+                "len": len,
+                "int": int,
+                "float": float,
+                "str": str,
+                "list": list,
+                "dict": dict,
+                "sum": sum,
+                "min": min,
+                "max": max,
+                "abs": abs,
+            }
         }
-    }
-
-    stdout_buffer = io.StringIO()
-
-    try:
-        signal.signal(signal.SIGALRM, timeout_handler)
-        signal.alarm(timeout)
-
-        with contextlib.redirect_stdout(stdout_buffer):
-            exec(code, safe_globals, {})
+        try:
+            signal.signal(signal.SIGALRM, timeout_handler)
+            signal.alarm(5)  # 5 second timeout
 
-        signal.alarm(0)
-        return {"status": "success", "output": stdout_buffer.getvalue()}
+            with contextlib.redirect_stdout(stdout_buffer):
+                exec(code, safe_globals, {})
 
-    except TimeoutException as e:
-        return {"status": "error", "output": str(e)}
+            signal.alarm(0)
+            return stdout_buffer.getvalue()
 
-    except Exception:
-        return {"status": "error", "output": traceback.format_exc()}
+        except TimeoutException as e:
+            return f"Error: {str(e)}"
 
+        except Exception:
+            return traceback.format_exc()
 
-# ---- THIS PART REGISTERS THE TOOL ----
-def python_runner(code: str) -> Dict[str, Any]:
-    """
-    Executes safe Python code and returns stdout or errors.
-    """
-    return run_python(code)