File size: 6,281 Bytes
d06fa57 | 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 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 | # Copyright (c) 2026-present, Yumeow. Licensed under the MIT License.
"""Execute data-only Python programs in an operating-system sandbox."""
from __future__ import annotations
import io
import json
from typing import Any, Dict
from ..runtime import get_sandbox_runner
from ..utils import bounded_value
from .base_tool import BaseTool, ToolMetadata
class LimitedWriter(io.StringIO):
"""String buffer retained for bounded-output shell utilities."""
SUFFIX = "...[truncated]"
def __init__(self, limit: int):
super().__init__()
self.limit = limit
def write(self, text: str) -> int:
remaining = self.limit - self.tell() - len(self.SUFFIX)
if remaining <= 0:
return len(text)
if len(text) > remaining:
super().write(text[:remaining])
super().write(self.SUFFIX)
return len(text)
return super().write(text)
@BaseTool.register("code_executor")
class CodeExecutorTool(BaseTool):
"""Implementation of the code executor tool."""
metadata = ToolMetadata(name="code_executor")
DEFAULT_TIMEOUT_SECONDS = 30
DEFAULT_MEMORY_LIMIT_MB = 1024
DEFAULT_OUTPUT_LIMIT_BYTES = 64 * 1024
MAX_TIMEOUT_SECONDS = 120
MAX_MEMORY_LIMIT_MB = 4096
MAX_OUTPUT_LIMIT_BYTES = 1024 * 1024
def execute(
self,
program: str,
timeout_seconds: int = DEFAULT_TIMEOUT_SECONDS,
memory_limit_mb: int = DEFAULT_MEMORY_LIMIT_MB,
output_limit_bytes: int = DEFAULT_OUTPUT_LIMIT_BYTES,
) -> Dict[str, Any]:
"""Execute Python code and return printed output.
1) Use `import sys, json; data_dict = json.loads(sys.stdin.read())` to
access data mapping variable names to lists.
2) Use `print()` to produce output.
3) The code is executed in an operating-system sandbox with resource limits.
4) Installed scientific Python libraries are available. Files outside the
ephemeral sandbox, networking, and child-process creation are unavailable.
Args:
program: Python code string to execute, starting with `import sys, json;
data_dict = json.loads(sys.stdin.read())` to access input data.
timeout_seconds: Wall-clock timeout in seconds. The effective value is capped.
memory_limit_mb: Address-space memory limit in MB. The effective value is capped.
output_limit_bytes: Limit on the amount of output (in bytes) that can be produced.
"""
timeout_seconds = self.bounded_int(
timeout_seconds, self.DEFAULT_TIMEOUT_SECONDS, 1, self.MAX_TIMEOUT_SECONDS
)
memory_limit_mb = self.bounded_int(
memory_limit_mb, self.DEFAULT_MEMORY_LIMIT_MB, 64, self.MAX_MEMORY_LIMIT_MB
)
output_limit_bytes = self.bounded_int(
output_limit_bytes, self.DEFAULT_OUTPUT_LIMIT_BYTES, 1024, self.MAX_OUTPUT_LIMIT_BYTES
)
stdin_text = json.dumps(self.serialization(self.context.data), ensure_ascii=False)
result = get_sandbox_runner().run(
self.extract_code(program),
stdin=stdin_text,
timeout_seconds=timeout_seconds,
memory_limit_mb=memory_limit_mb,
output_limit_bytes=output_limit_bytes,
interruption_event=getattr(self, "cancel_event", None),
)
if result.exit_code != 0:
raise RuntimeError(result.stderr or f"Sandboxed program exited with {result.exit_code}")
return {
"stdout": result.stdout,
"stderr": result.stderr,
"duration": result.duration_seconds,
"cpu_seconds": result.cpu_seconds,
"peak_memory_bytes": result.peak_memory_bytes,
"timeout_seconds": timeout_seconds,
"memory_limit_mb": memory_limit_mb,
"output_limit_bytes": output_limit_bytes,
}
@classmethod
def format_result_dict(cls, result: Dict[str, Any]) -> str:
"""Format a tool result for the language model.
Args:
result: Result mapping to format or update.
Returns:
str: The operation result.
"""
parts = []
if result["stdout"]:
parts.append(result["stdout"].rstrip())
if result["stderr"]:
parts.append(f"Program stderr:\n{result['stderr'].rstrip()}")
if not parts:
parts.append("Program completed successfully and printed no output.")
parts.append(f"(Execution duration: {result['duration']:.3f} seconds.)")
return "\n".join(parts)
@classmethod
def extract_code(cls, code: str) -> str:
"""Extract Python source from an optional Markdown fence."""
raw_code = str(code).strip()
if "```python" in raw_code:
return raw_code.split("```python")[-1].split("```")[0].strip()
if "```" not in raw_code:
return raw_code
parts = raw_code.split("```")
if len(parts) <= 1:
return raw_code
potential_code = parts[1]
if "\n" not in potential_code:
return potential_code.strip()
first_line, rest = potential_code.split("\n", 1)
return rest.strip() if first_line.strip().isalpha() else potential_code.strip()
@classmethod
def serialization(cls, value: Any) -> Any:
"""Convert context data to JSON-compatible values."""
if value is None or isinstance(value, (str, int, float, bool)):
return value
if isinstance(value, dict):
return {str(key): cls.serialization(item) for key, item in value.items()}
if isinstance(value, (list, tuple, set)):
return [cls.serialization(item) for item in value]
if hasattr(value, "tolist"):
return cls.serialization(value.tolist())
if hasattr(value, "item"):
return cls.serialization(value.item())
raise TypeError(f"unsupported data type {type(value).__name__}")
@classmethod
def bounded_int(cls, value: Any, default: int, minimum: int, maximum: int) -> int:
"""Return a bounded integer tool argument."""
return bounded_value(value, min=minimum, max=maximum, default=default, converter=int)
|