tools/python_tool.py

from __future__ import annotations

import ast
import math
import operator
import re
from typing import Any

_BINARY = {
    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 = {ast.UAdd: operator.pos, ast.USub: operator.neg}
_FUNCTIONS = {"abs": abs, "round": round, "sqrt": math.sqrt, "ceil": math.ceil, "floor": math.floor}


def calculate_expression(expression: str) -> dict[str, Any]:
    cleaned = _extract_expression(expression)
    value = _evaluate(ast.parse(cleaned, mode="eval").body)
    return {
        "ok": True,
        "source": "safe_calculator",
        "content": str(value),
        "metadata": {"expression": cleaned},
    }


def _extract_expression(text: str) -> str:
    text = text.strip().replace("×", "*").replace("÷", "/").replace("^", "**")
    match = re.search(r"(?:expression\s*[:=]\s*)?([0-9eE\.\s+\-*/%(),]+)$", text)
    candidate = (match.group(1) if match else text).strip()
    if not candidate:
        raise ValueError("No arithmetic expression found.")
    return candidate


def _evaluate(node: ast.AST) -> int | float:
    if isinstance(node, ast.Constant) and isinstance(node.value, (int, float)):
        return node.value
    if isinstance(node, ast.BinOp) and type(node.op) in _BINARY:
        return _BINARY[type(node.op)](_evaluate(node.left), _evaluate(node.right))
    if isinstance(node, ast.UnaryOp) and type(node.op) in _UNARY:
        return _UNARY[type(node.op)](_evaluate(node.operand))
    if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id in _FUNCTIONS:
        return _FUNCTIONS[node.func.id](*[_evaluate(arg) for arg in node.args])
    raise ValueError(f"Unsupported expression element: {ast.dump(node, include_attributes=False)}")