"""calculator_tool_demo.py — calculator @tool + create_agent"""

from dotenv import load_dotenv
from langchain.agents import create_agent
from langchain.tools import tool

load_dotenv()

ALLOWED = set("0123456789+-*/(). ")


def run_expression(expression: str) -> str:
    expr = expression.strip()
    if not expr or not all(ch in ALLOWED for ch in expr):
        return "Only numbers and + - * / ( ) allowed"
    try:
        value = eval(expr, {"__builtins__": {}}, {})
        return str(value)
    except Exception as exc:
        return f"Could not evaluate: {exc}"


@tool("calculator")
def calculator(expression: str) -> str:
    """Evaluate a math expression with +, -, *, / and parentheses."""
    return run_expression(expression)


agent = create_agent(model="openai:gpt-4o-mini", tools=[calculator])

questions = [
    "What is 24 times 17? Use the calculator tool.",
    "What is (100 + 15) divided by 5? Use calculator.",
]

for question in questions:
    print(f"\nQuestion: {question}")
    result = agent.invoke({"messages": [{"role": "user", "content": question}]})
    for msg in result["messages"]:
        if getattr(msg, "type", None) == "tool":
            print(f"calculator returned: {msg.content}")
    print(f"Answer: {result['messages'][-1].content}")