| """Restricted stateful Python for Nemotron Cascade-2 math tool calls. |
| |
| Tercet-R saw `stateful_python_code_exec` on `math/math_tool.jsonl` during |
| the Cascade-2 SFT stage. Observations are the Jupyter-style stdout / last |
| expression value, not a JSON envelope. |
| |
| The worker process only allows a math-oriented import whitelist and rejects |
| dunder access, so a tool call cannot read the filesystem or start a shell. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import ast |
| import io |
| import itertools |
| import json |
| import math |
| import os |
| import re |
| import statistics |
| import subprocess |
| import sys |
| import threading |
| from collections.abc import Iterable, Iterator, Mapping |
| from pathlib import Path |
| from queue import Empty, Queue |
| from typing import Any |
|
|
|
|
| NEMOTRON_PYTHON_EXEC_TOOL: dict[str, Any] = { |
| "type": "function", |
| "function": { |
| "name": "stateful_python_code_exec", |
| "description": ( |
| "Call this function to execute Python code in a stateful Jupyter " |
| "notebook environment. Python will respond with the output of the " |
| "execution or time out after 120.0 seconds." |
| ), |
| "parameters": { |
| "type": "object", |
| "properties": { |
| "code": { |
| "type": "string", |
| "description": "Code to execute", |
| } |
| }, |
| "required": ["code"], |
| }, |
| }, |
| } |
|
|
| CALCULATOR_TOOL: dict[str, Any] = { |
| "type": "function", |
| "function": { |
| "name": "calculator", |
| "description": "Evaluate a math expression and return the result.", |
| "parameters": { |
| "type": "object", |
| "properties": { |
| "expression": { |
| "type": "string", |
| "description": ( |
| "A math expression in standard calculator notation, for example " |
| "(3 + 4)^2 * sqrt(2) / ln(10). Use + - * / ^ (power), % (remainder), " |
| "! (factorial) and parentheses; the constants pi and e; and functions " |
| "such as sqrt, cbrt, exp, ln, log (natural, or log(x, base)), log10, " |
| "log2, sin, cos, tan, asin, acos, atan (radians), abs, round, floor, " |
| "ceil, factorial, comb, perm, gcd, lcm, min, max, sum, mean and median." |
| ), |
| } |
| }, |
| "required": ["expression"], |
| }, |
| }, |
| } |
|
|
| PYTHON_EXEC_TOOL_NAMES = frozenset( |
| { |
| "stateful_python_code_exec", |
| "python_code_exec", |
| "code_interpreter", |
| "python_exec", |
| "python", |
| } |
| ) |
| CALCULATOR_TOOL_NAMES = frozenset({"calculator", "calc"}) |
| DEFAULT_EXEC_TIMEOUT_SECONDS = 30.0 |
| _SRC_ROOT = Path(__file__).resolve().parents[1] |
|
|
| _session_lock = threading.Lock() |
| _default_session: StatefulPythonSession | None = None |
|
|
|
|
| def _normalize_tool_name(name: str) -> str: |
| return name.strip().lower().replace("-", "_") |
|
|
|
|
| def is_python_exec_tool_name(name: str) -> bool: |
| return _normalize_tool_name(name) in PYTHON_EXEC_TOOL_NAMES |
|
|
|
|
| def is_calculator_tool_name(name: str) -> bool: |
| return _normalize_tool_name(name) in CALCULATOR_TOOL_NAMES |
|
|
|
|
| def is_auto_math_tool_name(name: str) -> bool: |
| return is_python_exec_tool_name(name) or is_calculator_tool_name(name) |
|
|
|
|
| def is_calculator_expression(source: str) -> bool: |
| """A one-line calculator call runs in the calculator; several statements run as notebook code.""" |
| return "\n" not in source.strip() and ";" not in source |
|
|
|
|
| def code_from_arguments(arguments: Mapping[str, Any] | None) -> str: |
| if not arguments: |
| return "" |
| for key in ("code", "expression", "expr", "source"): |
| raw = arguments.get(key) |
| if isinstance(raw, str) and raw.strip(): |
| return raw |
| if raw is not None and key != "code": |
| return str(raw) |
| return "" |
|
|
|
|
| class CodeExecError(ValueError): |
| """Rejected or failed tool code.""" |
|
|
|
|
| class CalculatorError(CodeExecError): |
| """An expression the calculator cannot evaluate. The message says how to fix it.""" |
|
|
|
|
| class StatefulPythonSession: |
| """One long-lived restricted interpreter, matching the training tool.""" |
|
|
| def __init__(self, *, timeout_seconds: float = DEFAULT_EXEC_TIMEOUT_SECONDS) -> None: |
| self.timeout_seconds = timeout_seconds |
| self._lock = threading.Lock() |
| self._process: subprocess.Popen[str] | None = None |
|
|
| def close(self) -> None: |
| with self._lock: |
| self._kill_locked() |
|
|
| def reset(self) -> None: |
| with self._lock: |
| self._kill_locked() |
|
|
| def run(self, code: str) -> str: |
| text = code.strip() |
| if not text: |
| raise CodeExecError("code is empty") |
| with self._lock: |
| return self._run_locked(text) |
|
|
| def _run_locked(self, code: str) -> str: |
| process = self._ensure_process_locked() |
| try: |
| process.stdin.write(json.dumps({"code": code}, ensure_ascii=False) + "\n") |
| process.stdin.flush() |
| line = self._readline_locked(process, self.timeout_seconds) |
| except (BrokenPipeError, OSError) as error: |
| self._kill_locked() |
| raise CodeExecError(f"python worker died: {error}") from error |
| if not line: |
| stderr = "" |
| if process.stderr is not None: |
| try: |
| stderr = process.stderr.read() |
| except OSError: |
| stderr = "" |
| self._kill_locked() |
| detail = stderr.strip() or "python worker closed stdout" |
| raise CodeExecError(detail) |
| try: |
| payload = json.loads(line) |
| except json.JSONDecodeError as error: |
| self._kill_locked() |
| raise CodeExecError(f"python worker returned invalid JSON: {line!r}") from error |
| if not isinstance(payload, dict): |
| raise CodeExecError("python worker returned a non-object") |
| output = str(payload.get("output") or "") |
| if payload.get("ok") is True: |
| return output if output else "None" |
| raise CodeExecError(output or "execution failed") |
|
|
| def _readline_locked(self, process: subprocess.Popen[str], timeout: float) -> str: |
| if process.stdout is None: |
| raise CodeExecError("python worker has no stdout") |
| lines: Queue[str] = Queue() |
|
|
| def _read() -> None: |
| lines.put(process.stdout.readline() if process.stdout is not None else "") |
|
|
| reader = threading.Thread(target=_read, daemon=True) |
| reader.start() |
| reader.join(timeout) |
| if reader.is_alive(): |
| self._kill_locked() |
| raise CodeExecError(f"timed out after {timeout:g}s") |
| try: |
| return lines.get_nowait() |
| except Empty: |
| return "" |
|
|
| def _ensure_process_locked(self) -> subprocess.Popen[str]: |
| process = self._process |
| if process is not None and process.poll() is None: |
| return process |
| env = os.environ.copy() |
| pythonpath = env.get("PYTHONPATH", "") |
| env["PYTHONPATH"] = ( |
| str(_SRC_ROOT) if not pythonpath else f"{_SRC_ROOT}{os.pathsep}{pythonpath}" |
| ) |
| self._process = subprocess.Popen( |
| [sys.executable, "-m", "tiny_gdn.code_exec"], |
| stdin=subprocess.PIPE, |
| stdout=subprocess.PIPE, |
| stderr=subprocess.PIPE, |
| text=True, |
| bufsize=1, |
| env=env, |
| ) |
| if self._process.stdin is None or self._process.stdout is None: |
| self._kill_locked() |
| raise CodeExecError("failed to start python worker") |
| return self._process |
|
|
| def _kill_locked(self) -> None: |
| process = self._process |
| self._process = None |
| if process is None: |
| return |
| try: |
| process.kill() |
| process.wait(timeout=2) |
| except (OSError, subprocess.TimeoutExpired): |
| pass |
|
|
|
|
| def default_python_session() -> StatefulPythonSession: |
| global _default_session |
| with _session_lock: |
| if _default_session is None: |
| _default_session = StatefulPythonSession() |
| return _default_session |
|
|
|
|
| def reset_default_python_session() -> None: |
| global _default_session |
| with _session_lock: |
| session = _default_session |
| _default_session = None |
| if session is not None: |
| session.close() |
|
|
|
|
| MAX_EXPRESSION_CHARS = 10_000 |
| |
| _MAX_EXACT_POWER_BITS = 100_000 |
| _MAX_FACTORIAL = 100_000 |
| _MAX_LISTED_ITEMS = 1_000 |
| |
| _MAX_EXACT_INTEGER_BITS = 12_000 |
| _POWER = "_calculator_power" |
| _SYMBOLS = str.maketrans( |
| {"×": "*", "·": "*", "÷": "/", "−": "-", "–": "-", "π": "pi", "∞": "inf", "²": "**2", "³": "**3"} |
| ) |
|
|
|
|
| def calculate(expression: str, variables: Mapping[str, Any] | None = None) -> str: |
| """Evaluate one expression in calculator notation and format its result. |
| |
| `^` is a power and `n!` a factorial, as on a calculator, and functions |
| such as sqrt, ln and sin need no `math.` prefix. Python's `**`, `math.` |
| names and comprehensions work as well. `variables` are names that a |
| notebook in the same conversation defined. |
| """ |
| text = _calculator_text(expression) |
| try: |
| tree = ast.parse(text, mode="eval") |
| except SyntaxError as error: |
| raise CalculatorError( |
| f"could not read the expression ({error.msg}). " |
| "Write every operator out, for example 2*pi or 3*(4 + 5)." |
| ) from None |
| _reject_unsafe_ast(tree) |
| tree = ast.fix_missing_locations(_GuardedPower().visit(tree)) |
| |
| namespace: dict[str, Any] = {"__builtins__": {}, **_calculator_functions(), **(variables or {}), _POWER: _power} |
| try: |
| value = eval(compile(tree, "<calculator>", "eval"), namespace) |
| return _calculator_result(value) |
| except CodeExecError: |
| raise |
| except NameError as error: |
| raise CalculatorError( |
| f"unknown name {error.name or str(error)!r}. Call functions by name, such as sqrt(2), " |
| "ln(5) or sin(pi/6), and use the constants pi and e." |
| ) from None |
| except ZeroDivisionError: |
| raise CalculatorError("division by zero") from None |
| except OverflowError: |
| raise CalculatorError("the result is too large") from None |
| except ValueError as error: |
| if "domain" in str(error): |
| raise CalculatorError( |
| "math domain error: an input is outside the function's domain, " |
| "such as the square root or log of a negative number" |
| ) from None |
| raise CalculatorError(str(error)) from None |
| except Exception as error: |
| raise CalculatorError(str(error) or type(error).__name__) from None |
|
|
|
|
| def execute_calculator(expression: str) -> str: |
| return calculate(expression) |
|
|
|
|
| def execute_math_tool(name: str, arguments: Mapping[str, Any] | None) -> str: |
| source = code_from_arguments(arguments) |
| if is_calculator_tool_name(name) and is_calculator_expression(source): |
| return calculate(source) |
| return default_python_session().run(source) |
|
|
|
|
| def _calculator_text(expression: str) -> str: |
| """Calculator notation as Python: ^ as **, n! as factorial(n), √, π, × and ÷ spelled out.""" |
| text = expression.strip().translate(_SYMBOLS) |
| if text.endswith("=") and not text.endswith("=="): |
| text = text[:-1].rstrip() |
| if not text: |
| raise CalculatorError("the expression is empty") |
| if len(text) > MAX_EXPRESSION_CHARS: |
| raise CalculatorError(f"the expression is longer than {MAX_EXPRESSION_CHARS:,} characters") |
| text = re.sub(r"√\s*(?=\()", "sqrt", text) |
| text = re.sub(r"√\s*(\d+(?:\.\d*)?|[A-Za-z_]\w*)", r"sqrt(\1)", text) |
| text = re.sub(r"\bmod\b", "%", text) |
| return _expand_factorials(text.replace("^", "**")) |
|
|
|
|
| def _expand_factorials(text: str) -> str: |
| """Rewrite postfix factorials, 5! or (n + 1)!, as factorial() calls. != stays a comparison.""" |
| position = 0 |
| while (index := text.find("!", position)) >= 0: |
| if text.startswith("!=", index): |
| position = index + 2 |
| continue |
| start = _operand_start(text, index) |
| if start is None: |
| raise CalculatorError("! needs a number, a name or a parenthesized expression before it, as in 5! or (n + 1)!") |
| text = f"{text[:start]}factorial({text[start:index].strip()}){text[index + 1:]}" |
| position = start |
| return text |
|
|
|
|
| def _operand_start(text: str, end: int) -> int | None: |
| """Where the operand that ends just before `end` starts: a number, a name, a call or a parenthesized group.""" |
| index = end |
| while index > 0 and text[index - 1].isspace(): |
| index -= 1 |
| stop = index |
| if index > 0 and text[index - 1] == ")": |
| depth = 0 |
| cursor = index - 1 |
| while cursor >= 0: |
| if text[cursor] == ")": |
| depth += 1 |
| elif text[cursor] == "(": |
| depth -= 1 |
| if depth == 0: |
| break |
| cursor -= 1 |
| if cursor < 0: |
| return None |
| index = cursor |
| while index > 0 and (text[index - 1].isalnum() or text[index - 1] in "_."): |
| index -= 1 |
| return index if index < stop else None |
|
|
|
|
| class _GuardedPower(ast.NodeTransformer): |
| """Send every power through `_power`, which keeps huge exact integers from stalling the worker.""" |
|
|
| def visit_BinOp(self, node: ast.BinOp) -> ast.AST: |
| self.generic_visit(node) |
| if isinstance(node.op, ast.Pow): |
| return ast.Call(func=ast.Name(id=_POWER, ctx=ast.Load()), args=[node.left, node.right], keywords=[]) |
| return node |
|
|
|
|
| def _power(base: Any, exponent: Any, modulus: Any = None) -> Any: |
| if modulus is not None: |
| return pow(base, exponent, modulus) |
| if ( |
| type(base) is int |
| and type(exponent) is int |
| and exponent > 0 |
| and abs(base) > 1 |
| and exponent * math.log2(abs(base)) > _MAX_EXACT_POWER_BITS |
| ): |
| return float(base) ** exponent |
| return base**exponent |
|
|
|
|
| def _factorial(value: Any) -> Any: |
| if isinstance(value, float) and value.is_integer(): |
| value = int(value) |
| if isinstance(value, int) and not isinstance(value, bool): |
| if value < 0: |
| raise CalculatorError("factorial is not defined for negative integers") |
| if value > _MAX_FACTORIAL: |
| raise CalculatorError(f"factorial({value}) is too large; lgamma({value} + 1) gives its natural log") |
| return math.factorial(value) |
| return math.gamma(value + 1) |
|
|
|
|
| def _values(arguments: tuple[Any, ...]) -> list[Any]: |
| """An aggregate's numbers, given separately or as one list: mean(1, 2, 3) or mean([1, 2, 3]).""" |
| if len(arguments) == 1 and isinstance(arguments[0], Iterable) and not isinstance(arguments[0], (str, bytes)): |
| return list(arguments[0]) |
| return list(arguments) |
|
|
|
|
| def _root(value: Any, degree: Any) -> Any: |
| """The real nth root, so root(-8, 3) is -2.""" |
| if isinstance(value, (int, float)) and value < 0 and float(degree).is_integer() and int(degree) % 2 == 1: |
| return -((-value) ** (1 / degree)) |
| return value ** (1 / degree) |
|
|
|
|
| def _calculator_functions() -> dict[str, Any]: |
| names: dict[str, Any] = {key: getattr(math, key) for key in dir(math) if not key.startswith("_")} |
| names.update( |
| { |
| "ln": lambda value: math.log(value), |
| "arcsin": math.asin, |
| "arccos": math.acos, |
| "arctan": math.atan, |
| "arctan2": math.atan2, |
| "sec": lambda value: 1 / math.cos(value), |
| "csc": lambda value: 1 / math.sin(value), |
| "cot": lambda value: 1 / math.tan(value), |
| "sign": lambda value: (value > 0) - (value < 0), |
| "sgn": lambda value: (value > 0) - (value < 0), |
| "root": _root, |
| "factorial": _factorial, |
| "fact": _factorial, |
| "nCr": math.comb, |
| "choose": math.comb, |
| "nPr": math.perm, |
| "pow": _power, |
| "abs": abs, |
| "round": round, |
| "int": int, |
| "float": float, |
| "divmod": divmod, |
| "range": range, |
| "len": len, |
| "sorted": sorted, |
| "list": list, |
| "tuple": tuple, |
| "sum": lambda *values: sum(_values(values)), |
| "min": lambda *values: min(_values(values)), |
| "max": lambda *values: max(_values(values)), |
| "prod": lambda *values: math.prod(_values(values)), |
| "product": lambda *values: math.prod(_values(values)), |
| "mean": lambda *values: statistics.fmean(_values(values)), |
| "avg": lambda *values: statistics.fmean(_values(values)), |
| "average": lambda *values: statistics.fmean(_values(values)), |
| "median": lambda *values: statistics.median(_values(values)), |
| "mode": lambda *values: statistics.mode(_values(values)), |
| "stdev": lambda *values: statistics.stdev(_values(values)), |
| "std": lambda *values: statistics.stdev(_values(values)), |
| "pstdev": lambda *values: statistics.pstdev(_values(values)), |
| "variance": lambda *values: statistics.variance(_values(values)), |
| "var": lambda *values: statistics.variance(_values(values)), |
| "pvariance": lambda *values: statistics.pvariance(_values(values)), |
| "phi": (1 + math.sqrt(5)) / 2, |
| "deg": math.pi / 180, |
| "math": math, |
| } |
| ) |
| return names |
|
|
|
|
| def _calculator_result(value: Any) -> str: |
| if value is None or isinstance(value, bool): |
| return str(value) |
| if isinstance(value, int): |
| return _exact_integer(value) |
| if isinstance(value, float): |
| if math.isnan(value): |
| return "nan" |
| if math.isinf(value): |
| return "inf" if value > 0 else "-inf" |
| if value.is_integer() and abs(value) < 1e15: |
| return str(int(value)) |
| return format(value, ".12g") |
| if isinstance(value, complex): |
| if value.imag == 0: |
| return _calculator_result(value.real) |
| if value.real == 0: |
| return f"{_calculator_result(value.imag)}i" |
| sign = "-" if value.imag < 0 else "+" |
| return f"{_calculator_result(value.real)} {sign} {_calculator_result(abs(value.imag))}i" |
| if isinstance(value, (list, tuple, set, frozenset, range, Iterator)): |
| items = list(itertools.islice(iter(value), _MAX_LISTED_ITEMS + 1)) |
| shown = ", ".join(_calculator_result(item) for item in items[:_MAX_LISTED_ITEMS]) |
| if len(items) > _MAX_LISTED_ITEMS: |
| hidden = f"{len(value) - _MAX_LISTED_ITEMS:,} more" if isinstance(value, (list, tuple, set, frozenset, range)) else "more" |
| shown = f"{shown}, … ({hidden})" |
| return f"({shown})" if isinstance(value, tuple) else f"[{shown}]" |
| return str(value) |
|
|
|
|
| def _exact_integer(value: int) -> str: |
| if value.bit_length() <= _MAX_EXACT_INTEGER_BITS: |
| return str(value) |
| magnitude = abs(value) |
| shift = max(int(math.log10(magnitude)) - 12, 0) |
| head = str(magnitude // 10**shift) |
| digits = len(head) + shift |
| mantissa = f"{head[0]}.{head[1:13]}".rstrip("0").rstrip(".") |
| sign = "-" if value < 0 else "" |
| return f"{sign}{mantissa}e+{digits - 1} (an integer with {digits:,} digits)" |
|
|
|
|
| def _reject_unsafe_ast(tree: ast.AST) -> None: |
| for node in ast.walk(tree): |
| if isinstance(node, (ast.ClassDef, ast.AsyncFunctionDef)): |
| raise CodeExecError("class definitions are blocked") |
| if isinstance(node, ast.Attribute) and node.attr.startswith("_"): |
| raise CodeExecError("dunder attribute access is blocked") |
| if isinstance(node, ast.Name) and node.id.startswith("_"): |
| raise CodeExecError("dunder names are blocked") |
| if isinstance(node, ast.Call) and isinstance(node.func, ast.Name): |
| if node.func.id in _BANNED_CALLS: |
| raise CodeExecError(f"{node.func.id}() is blocked") |
|
|
|
|
| _BANNED_CALLS = frozenset( |
| { |
| "eval", |
| "exec", |
| "compile", |
| "open", |
| "input", |
| "breakpoint", |
| "getattr", |
| "setattr", |
| "delattr", |
| "globals", |
| "locals", |
| "vars", |
| "dir", |
| "help", |
| "__import__", |
| "memoryview", |
| "exit", |
| "quit", |
| } |
| ) |
|
|
|
|
| def _format_result(value: Any) -> str: |
| if value is None: |
| return "None" |
| if isinstance(value, bool): |
| return str(value) |
| if isinstance(value, float): |
| if value.is_integer() and abs(value) < 1e15: |
| return str(int(value)) |
| return format(value, ".12g") |
| return str(value) |
|
|
|
|
| _ALLOWED_IMPORT_ROOTS = frozenset( |
| { |
| "cmath", |
| "collections", |
| "copy", |
| "decimal", |
| "fractions", |
| "functools", |
| "itertools", |
| "json", |
| "math", |
| "mpmath", |
| "numbers", |
| "numpy", |
| "operator", |
| "re", |
| "statistics", |
| "string", |
| "sympy", |
| "textwrap", |
| "unicodedata", |
| } |
| ) |
|
|
| _WORKER_NAMESPACE: dict[str, Any] | None = None |
|
|
|
|
| def _allowed_import( |
| name: str, |
| globals: dict[str, Any] | None = None, |
| locals: dict[str, Any] | None = None, |
| fromlist: tuple[str, ...] = (), |
| level: int = 0, |
| ) -> Any: |
| root = name.split(".")[0] |
| if root not in _ALLOWED_IMPORT_ROOTS: |
| raise ImportError(f"import of {name!r} is blocked") |
| return __import__(name, globals, locals, fromlist, level) |
|
|
|
|
| def _worker_namespace() -> dict[str, Any]: |
| global _WORKER_NAMESPACE |
| if _WORKER_NAMESPACE is None: |
| builtins = { |
| "abs": abs, |
| "all": all, |
| "any": any, |
| "bin": bin, |
| "bool": bool, |
| "bytes": bytes, |
| "chr": chr, |
| "complex": complex, |
| "dict": dict, |
| "divmod": divmod, |
| "enumerate": enumerate, |
| "filter": filter, |
| "float": float, |
| "format": format, |
| "frozenset": frozenset, |
| "hex": hex, |
| "int": int, |
| "isinstance": isinstance, |
| "issubclass": issubclass, |
| "iter": iter, |
| "len": len, |
| "list": list, |
| "map": map, |
| "max": max, |
| "min": min, |
| "next": next, |
| "oct": oct, |
| "ord": ord, |
| "pow": pow, |
| "print": print, |
| "range": range, |
| "repr": repr, |
| "reversed": reversed, |
| "round": round, |
| "set": set, |
| "slice": slice, |
| "sorted": sorted, |
| "str": str, |
| "sum": sum, |
| "tuple": tuple, |
| "zip": zip, |
| "True": True, |
| "False": False, |
| "None": None, |
| "__import__": _allowed_import, |
| } |
| _WORKER_NAMESPACE = { |
| "__builtins__": builtins, |
| "__name__": "__tool__", |
| "math": math, |
| } |
| return _WORKER_NAMESPACE |
|
|
|
|
| def run_cell(code: str) -> str: |
| tree = ast.parse(code) |
| _reject_unsafe_import_roots(tree) |
| _reject_unsafe_ast(tree) |
| namespace = _worker_namespace() |
| buffer = io.StringIO() |
| previous = sys.stdout |
| sys.stdout = buffer |
| try: |
| if tree.body and isinstance(tree.body[-1], ast.Expr): |
| body = tree.body[:-1] |
| last = tree.body[-1] |
| if body: |
| exec( |
| compile(ast.Module(body, type_ignores=[]), "<tool>", "exec"), |
| namespace, |
| namespace, |
| ) |
| value = eval( |
| compile(ast.Expression(last.value), "<tool>", "eval"), |
| namespace, |
| namespace, |
| ) |
| printed = buffer.getvalue() |
| if value is None: |
| return printed if printed else "None" |
| rendered = _format_result(value) |
| return f"{printed}{rendered}" if printed else rendered |
| exec(compile(tree, "<tool>", "exec"), namespace, namespace) |
| printed = buffer.getvalue() |
| return printed if printed else "None" |
| finally: |
| sys.stdout = previous |
|
|
|
|
| def notebook_variables() -> dict[str, Any]: |
| """Names this worker's notebook defined, which calculator calls in the same conversation can use.""" |
| return {key: value for key, value in _worker_namespace().items() if not key.startswith("_") and key != "math"} |
|
|
|
|
| def _reject_unsafe_import_roots(tree: ast.AST) -> None: |
| for node in ast.walk(tree): |
| if isinstance(node, ast.Import): |
| for alias in node.names: |
| root = alias.name.split(".")[0] |
| if root not in _ALLOWED_IMPORT_ROOTS: |
| raise CodeExecError(f"import of {alias.name!r} is blocked") |
| elif isinstance(node, ast.ImportFrom): |
| root = (node.module or "").split(".")[0] |
| if root not in _ALLOWED_IMPORT_ROOTS: |
| raise CodeExecError(f"import of {node.module!r} is blocked") |
|
|
|
|
| def _worker_loop() -> None: |
| for line in sys.stdin: |
| line = line.strip() |
| if not line: |
| continue |
| try: |
| request = json.loads(line) |
| if "expression" in request: |
| output = calculate(str(request.get("expression") or ""), notebook_variables()) |
| else: |
| output = run_cell(str(request.get("code") or "")) |
| sys.stdout.write(json.dumps({"ok": True, "output": output}, ensure_ascii=False) + "\n") |
| except Exception as error: |
| detail = str(error) or type(error).__name__ |
| sys.stdout.write( |
| json.dumps({"ok": False, "output": detail}, ensure_ascii=False) + "\n" |
| ) |
| sys.stdout.flush() |
|
|
|
|
| if __name__ == "__main__": |
| _worker_loop() |
|
|
