calculator.py•4.48 kB
"""
Calculator Tool
Author: Yobie Benjamin
Version: 0.9
Date: August 1, 2025
A tool for performing mathematical calculations.
Supports basic arithmetic, advanced functions, and symbolic math.
"""
import ast
import math
import operator as op
from typing import Any
from pydantic import BaseModel, Field
from .base import BaseTool, ToolResult
# Safe operators for evaluation
SAFE_OPS = {
ast.Add: op.add,
ast.Sub: op.sub,
ast.Mult: op.mul,
ast.Div: op.truediv,
ast.Pow: op.pow,
ast.Mod: op.mod,
ast.FloorDiv: op.floordiv,
ast.USub: op.neg,
ast.UAdd: op.pos,
}
# Safe functions
SAFE_FUNCS = {
'abs': abs,
'round': round,
'min': min,
'max': max,
'sum': sum,
'len': len,
'sin': math.sin,
'cos': math.cos,
'tan': math.tan,
'sqrt': math.sqrt,
'log': math.log,
'log10': math.log10,
'exp': math.exp,
'pi': math.pi,
'e': math.e,
}
class CalculatorParams(BaseModel):
"""Parameters for the calculator tool."""
expression: str = Field(
...,
description="Mathematical expression to evaluate (e.g., '2 + 2', 'sqrt(16)', 'pi * r**2')"
)
class CalculatorTool(BaseTool):
"""
Calculator tool for mathematical computations.
This tool safely evaluates mathematical expressions without
allowing arbitrary code execution.
"""
@property
def name(self) -> str:
return "calculator"
@property
def description(self) -> str:
return "Perform mathematical calculations safely"
@property
def parameters(self) -> type[BaseModel]:
return CalculatorParams
async def execute(self, expression: str) -> ToolResult:
"""
Evaluate a mathematical expression.
Args:
expression: Mathematical expression to evaluate
Returns:
ToolResult with the calculated value
"""
try:
# Parse the expression
node = ast.parse(expression, mode='eval')
# Evaluate safely
result = self._eval_node(node.body)
return ToolResult(
success=True,
data={
"expression": expression,
"result": result,
"type": type(result).__name__
}
)
except ZeroDivisionError:
return ToolResult(
success=False,
error="Division by zero",
data={"expression": expression}
)
except Exception as e:
return ToolResult(
success=False,
error=f"Calculation error: {str(e)}",
data={"expression": expression}
)
def _eval_node(self, node: ast.AST) -> Any:
"""
Safely evaluate an AST node.
Args:
node: AST node to evaluate
Returns:
Evaluated result
Raises:
ValueError: If unsafe operation detected
"""
if isinstance(node, ast.Constant):
return node.value
elif isinstance(node, ast.Name):
# Allow only safe names
if node.id in SAFE_FUNCS:
return SAFE_FUNCS[node.id]
raise ValueError(f"Unsafe name: {node.id}")
elif isinstance(node, ast.UnaryOp):
op_type = type(node.op)
if op_type not in SAFE_OPS:
raise ValueError(f"Unsafe operator: {op_type}")
return SAFE_OPS[op_type](self._eval_node(node.operand))
elif isinstance(node, ast.BinOp):
op_type = type(node.op)
if op_type not in SAFE_OPS:
raise ValueError(f"Unsafe operator: {op_type}")
left = self._eval_node(node.left)
right = self._eval_node(node.right)
return SAFE_OPS[op_type](left, right)
elif isinstance(node, ast.Call):
# Allow only safe function calls
if isinstance(node.func, ast.Name):
func_name = node.func.id
if func_name in SAFE_FUNCS:
args = [self._eval_node(arg) for arg in node.args]
return SAFE_FUNCS[func_name](*args)
raise ValueError(f"Unsafe function call")
else:
raise ValueError(f"Unsafe node type: {type(node)}")