from smolagents import Tool

class PythonCalcTool(Tool):
    name = "python_calc"
    description = (
        "Can use a normal calculator for probability questions. If needed Execute Python code to compute numeric answers for calculation problems. Use tool to solve math, physics, or distance/time problems. This tool is used for **all mathematical calculations, complex logic processing, and numerical comparisons (e.g., comparing probabilities or finding maximum values)."
    )
    inputs = {
        "code": {"type": "string", "description": "Python code that sets a variable 'result' with the answer."}
    }
    output_type = "string"

    def forward(self, code: str) -> str:
        """
        Executes the code safely and returns the result as a string.
        Code must set a variable 'result' which is returned.
        """
        try:
            local_vars = {}
            exec(code, {}, local_vars)
            # The agent should always set 'result'
            answer = local_vars.get("result", None)
            if answer is None:
                return "NA"
            return str(answer)
        except Exception as e:
            return "NA"