vgtc-api / src /hermes /tools /python_executor.py
vora-sonnet's picture
Upload folder using huggingface_hub (part 2)
b9f94e1 verified
Raw
History Blame Contribute Delete
4.59 kB
"""Python execution tool implementation."""
from __future__ import annotations
import asyncio
import contextlib
import logging
import os
import sys
import tempfile
from pathlib import Path
from typing import Any
from hermes.tools.base.tool import BaseTool, ToolSchema
logger = logging.getLogger(__name__)
class PythonExecutionTool(BaseTool):
"""Tool for executing Python code with safety measures."""
def __init__(self, timeout: int = 30, max_output_size: int = 10000, max_memory_mb: int = 256) -> None:
super().__init__()
self.timeout = timeout
self.max_output_size = max_output_size
self.max_memory_mb = max_memory_mb
def _define_schema(self) -> ToolSchema:
return ToolSchema(
name="python_executor",
description="Execute Python code in an isolated environment with resource limits",
parameters={
"code": {
"type": "string",
"description": "Python code to execute",
},
"timeout": {
"type": "integer",
"description": "Execution timeout in seconds",
"default": 30,
},
},
required=["code"],
category="execution",
tags=["python", "execute", "code"],
)
async def execute(self, **kwargs: Any) -> dict[str, Any]:
"""Execute Python code with isolation."""
code = kwargs["code"]
timeout = kwargs.get("timeout", self.timeout)
restricted_code = self._inject_safety(code)
with tempfile.NamedTemporaryFile(
mode="w", suffix=".py", delete=False, encoding="utf-8"
) as f:
f.write(restricted_code)
temp_path = f.name
try:
env = os.environ.copy()
sensitive_keys = [k for k in env if any(word in k.upper() for word in
["API_KEY", "APIKEY", "SECRET", "TOKEN", "PASSWORD", "CREDENTIAL", "AUTH"])]
for k in sensitive_keys:
del env[k]
process = await asyncio.create_subprocess_exec(
sys.executable,
"-I",
temp_path,
stdout=asyncio.subprocess.PIPE,
stderr=asyncio.subprocess.PIPE,
env=env,
cwd=tempfile.gettempdir(),
)
try:
stdout, stderr = await asyncio.wait_for(
process.communicate(), timeout=timeout
)
except TimeoutError:
process.kill()
await process.wait()
return {
"success": False,
"error": f"Execution timed out after {timeout}s",
"stdout": "",
"stderr": "",
"return_code": -1,
}
stdout_str = stdout.decode("utf-8", errors="replace")[:self.max_output_size]
stderr_str = stderr.decode("utf-8", errors="replace")[:self.max_output_size]
return {
"success": process.returncode == 0,
"stdout": stdout_str,
"stderr": stderr_str,
"return_code": process.returncode,
}
except Exception as e:
logger.error(f"Python execution error: {e}")
return {
"success": False,
"error": str(e),
"stdout": "",
"stderr": "",
"return_code": -1,
}
finally:
with contextlib.suppress(Exception):
Path(temp_path).unlink(missing_ok=True)
def _inject_safety(self, code: str) -> str:
"""Inject safety restrictions into code — blocks imports and subprocess access."""
banner = """import sys as _sys, builtins as _builtins
# Block __import__ — prevents all module imports
_builtins.__import__ = None
# Nullify dangerous modules so even if __import__ is somehow recovered they fail
for _mod in ['os', 'subprocess', 'shutil', 'socket', 'ctypes', 'pickle',
'marshal', 'compile', 'codeop', 'code', 'inspect',
'importlib', 'pkgutil', 'mmap', 'webbrowser', 'antigravity',
'ctypes', '_ctypes', '_winapi', 'posix', 'grp', 'pwd']:
_sys.modules[_mod] = None
del _sys, _builtins
"""
return banner + code