from typing import Literal
from autogen_core import CancellationToken
from autogen_core.tools import FunctionTool, ImportFromModule


async def calculate(
    operation:  Literal["add", "subtract", "multiply", "divide"],
    x: float,
    y: float,
) -> float:
    """Perform basic arithmetic operations on two numbers."""
    if operation == "add":
        return x + y
    elif operation == "subtract":
        return x - y
    elif operation == "multiply":
        return x * y
    elif operation == "divide":
        if y == 0:
            raise ValueError("Cannot divide by zero")
        return x / y
    else:
        raise ValueError(f"Unknown operation: {operation}")


# Create the calculator tool
calculator_tool = FunctionTool(
    func=calculate,
    description="A calculator that can add, subtract, multiply, or divide two numbers.",
    global_imports=[
        ImportFromModule("typing", ("Literal",)),
        ImportFromModule("autogen_core", ("CancellationToken",)),
        ImportFromModule("autogen_core.tools", ("FunctionTool",))
    ]
)

tool_config = calculator_tool.dump_component()
print(tool_config.model_dump())
 

# load
loaded_calculator = FunctionTool.load_component(tool_config)