File size: 1,457 Bytes
e945848 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 | from typing import Any, Dict
from smolagents.tools import Tool
import io
import contextlib
import traceback
import signal
class TimeoutException(Exception):
pass
def timeout_handler(signum, frame):
raise TimeoutException("Execution timed out.")
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,
}
}
try:
signal.signal(signal.SIGALRM, timeout_handler)
signal.alarm(5) # 5 second timeout
with contextlib.redirect_stdout(stdout_buffer):
exec(code, safe_globals, {})
signal.alarm(0)
return stdout_buffer.getvalue()
except TimeoutException as e:
return f"Error: {str(e)}"
except Exception:
return traceback.format_exc()
|