from smolagents import CodeAgent, DuckDuckGoSearchTool, GradioUI, InferenceClientModel, Tool


class AddTool(Tool):
    name = "add"
    description = "Adds two numbers together and returns the result."
    inputs = {
        "a": {"type": "number", "description": "The first number"},
        "b": {"type": "number", "description": "The second number"},
    }
    output_type = "number"

    def forward(self, a: float, b: float) -> float:
        return a + b


class SubtractTool(Tool):
    name = "subtract"
    description = "Subtracts the second number from the first and returns the result."
    inputs = {
        "a": {"type": "number", "description": "The number to subtract from"},
        "b": {"type": "number", "description": "The number to subtract"},
    }
    output_type = "number"

    def forward(self, a: float, b: float) -> float:
        return a - b


agent = CodeAgent(
    tools=[AddTool(), SubtractTool(), DuckDuckGoSearchTool()],
    model=InferenceClientModel(model_id="Qwen/Qwen2.5-Coder-32B-Instruct"),
    name="calculator_agent",
    description="An agent that can perform arithmetic (addition, subtraction) and search the web.",
    max_steps=10,
)

if __name__ == "__main__":
    GradioUI(agent).launch()