Source code for mailos.tools.python_interpreter
"""Python code interpreter tool."""
import signal
import sys
import traceback
from io import StringIO
from typing import Dict
from mailos.utils.logger_utils import logger
from mailos.vendors.models import Tool
def timeout_handler(signum, frame):
"""Handle timeout signal."""
raise TimeoutError("Code execution timed out")
[docs]
def execute_python(code: str, timeout: int = 5) -> Dict:
"""Execute Python code and return the output.
Args:
code: Python code to execute
timeout: Maximum execution time in seconds (default: 5)
Returns:
Dict with execution status, output/error, and any printed output
"""
# Capture stdout and stderr
old_stdout = sys.stdout
old_stderr = sys.stderr
stdout = StringIO()
stderr = StringIO()
sys.stdout = stdout
sys.stderr = stderr
# Set up the timeout handler
old_handler = signal.signal(signal.SIGALRM, timeout_handler)
signal.alarm(timeout)
result = None
try:
# Split code into blocks that can be executed independently
lines = code.strip().split("\n")
current_block = []
exec_globals = {}
for line in lines:
current_block.append(line)
try:
# Try to compile the current block
block_code = "\n".join(current_block)
compiled_code = compile(block_code, "<string>", "exec")
# If compilation succeeds, execute the block
exec(compiled_code, exec_globals)
current_block = [] # Reset for next block
except SyntaxError:
# If we get a syntax error, keep accumulating lines
# This handles multi-line statements
continue
except Exception as e:
# For runtime errors, return immediately
error_type = type(e).__name__
error_msg = f"{error_type}: {str(e)}"
error_traceback = traceback.format_exc()
return {
"status": "error",
"error": error_msg,
"traceback": error_traceback,
"output": stdout.getvalue() + stderr.getvalue(),
}
# Try to compile and execute any remaining lines
if current_block:
try:
block_code = "\n".join(current_block)
compiled_code = compile(block_code, "<string>", "exec")
exec(compiled_code, exec_globals)
except Exception as e:
error_type = type(e).__name__
error_msg = f"{error_type}: {str(e)}"
error_traceback = traceback.format_exc()
return {
"status": "error",
"error": error_msg,
"traceback": error_traceback,
"output": stdout.getvalue() + stderr.getvalue(),
}
# Get any printed output
output = stdout.getvalue() + stderr.getvalue() # Include stderr in output
return {
"status": "success",
"output": output,
"result": str(result) if result is not None else None,
}
except TimeoutError as e:
error_msg = f"TimeoutError: {str(e)}"
logger.error(f"Code execution timed out: {error_msg}")
return {
"status": "error",
"error": error_msg,
"traceback": error_msg,
"output": stdout.getvalue() + stderr.getvalue(),
}
except Exception as e:
error_type = type(e).__name__
error_msg = f"{error_type}: {str(e)}"
error_traceback = traceback.format_exc()
logger.error(f"Error executing Python code: {error_traceback}")
return {
"status": "error",
"error": error_msg,
"traceback": error_traceback,
"output": stdout.getvalue() + stderr.getvalue(),
}
finally:
# Restore stdout and stderr
sys.stdout = old_stdout
sys.stderr = old_stderr
# Restore the old signal handler and cancel the alarm
signal.alarm(0)
signal.signal(signal.SIGALRM, old_handler)
# Define the Python interpreter tool
python_interpreter_tool = Tool(
name="execute_python",
description="Execute Python code and return the output",
parameters={
"type": "object",
"properties": {
"code": {"type": "string", "description": "Python code to execute"}
},
},
required_params=["code"],
function=execute_python,
)