Spaces:
Sleeping
Sleeping
File size: 3,657 Bytes
b2931f4 | 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 | """calculator(expression) — arithmetic on figures the model pulled from context.
Why this exists: LLMs are unreliable at multi-digit arithmetic (growth rates,
margins, sums across years). Far better to have the model *extract* the numbers
and delegate the math to real code.
Why not eval(): `eval("__import__('os').system('...')")` is remote code
execution. A figure could even arrive via a prompt-injected chunk. So we parse
to an AST and walk it, permitting ONLY numeric literals and arithmetic
operators — every other node type (names, calls, attributes, subscripts)
raises. This is an allowlist, not a blocklist: anything we didn't explicitly
permit is rejected by default.
"""
from __future__ import annotations
import ast
import operator
from typing import Any
# Allowlisted operators → their implementing functions. Anything not here
# (e.g. bitwise, matmul) is rejected.
_BIN_OPS: dict[type[ast.operator], Any] = {
ast.Add: operator.add,
ast.Sub: operator.sub,
ast.Mult: operator.mul,
ast.Div: operator.truediv,
ast.FloorDiv: operator.floordiv,
ast.Mod: operator.mod,
ast.Pow: operator.pow,
}
_UNARY_OPS: dict[type[ast.unaryop], Any] = {
ast.UAdd: operator.pos,
ast.USub: operator.neg,
}
# Guardrail: cap exponent magnitude so `10 ** 10**9` can't pin a CPU / OOM.
_MAX_EXPONENT = 1000
def _eval_node(node: ast.AST) -> float:
if isinstance(node, ast.Expression):
return _eval_node(node.body)
if isinstance(node, ast.Constant):
if isinstance(node.value, bool) or not isinstance(node.value, (int, float)):
raise ValueError(f"Only numeric literals allowed, got {node.value!r}")
return float(node.value)
if isinstance(node, ast.UnaryOp):
op = _UNARY_OPS.get(type(node.op))
if op is None:
raise ValueError(f"Operator {type(node.op).__name__} not allowed")
return op(_eval_node(node.operand))
if isinstance(node, ast.BinOp):
op = _BIN_OPS.get(type(node.op))
if op is None:
raise ValueError(f"Operator {type(node.op).__name__} not allowed")
left, right = _eval_node(node.left), _eval_node(node.right)
if isinstance(node.op, ast.Pow) and abs(right) > _MAX_EXPONENT:
raise ValueError(f"Exponent {right} exceeds limit {_MAX_EXPONENT}")
return op(left, right)
# Any other node — Name, Call, Attribute, Subscript, etc. — is rejected.
raise ValueError(f"Expression element {type(node).__name__} not allowed")
def calculator(expression: str) -> dict[str, Any]:
"""Evaluate an arithmetic `expression` and return the numeric result.
Supports + - * / // % ** and parentheses over numeric literals only.
Returns {"result": <float>} on success or {"error": <message>} on failure
— tools return errors as data (not exceptions) so the agent can read the
message and retry rather than crashing the graph.
"""
try:
tree = ast.parse(expression, mode="eval")
result = _eval_node(tree)
return {"result": result}
except ZeroDivisionError:
return {"error": "division by zero"}
except (ValueError, SyntaxError) as e:
return {"error": str(e)}
if __name__ == "__main__":
# Sanity: valid arithmetic, plus rejection of an injection attempt.
for expr in [
"(383285 - 394328) / 394328 * 100", # YoY % change
"200583 + 85200", # iPhone + Services
"2 ** 4000", # exponent guard
"__import__('os').system('echo pwned')", # must be rejected
]:
print(f"{expr!r:50s} -> {calculator(expr)}")
|