import os

HF_TOKEN = os.getenv("HF_TOKEN")
from smolagents import Tool, FinalAnswerTool, ToolCallingAgent, DuckDuckGoSearchTool, tool, LiteLLMModel, PythonInterpreterTool
import datetime
import requests
import pytz
from tools.final_answer import FinalAnswerTool
from Gradio_UI import GradioUI


@tool
def sum_tool(a: float, b: float) -> float:
    """
    Returns the sum of two numbers

    Args:
        a: The first value.
        b: The second value.

    """
    return a + b

@tool
def multiply_tool(a: float, b: float) -> float:
    """
    Returns the product of two numbers.

    Args:
        a (float): The first value.
        b (float): The second value.

    Returns:
        float: Product of a and b
    """
    return a * b

@tool
def sub_tool(a: float, b: float) -> float:
    """
    Returns the difference of two numbers

    Args:
        a: The first value.
        b: The second value.

    """
    return a - b


final_answer = FinalAnswerTool()

model = LiteLLMModel(
    model_id="huggingface/google/gemma-2-2b-it",
    api_key=HF_TOKEN
)

agent = ToolCallingAgent(
    tools=[sum_tool, sub_tool, multiply_tool],
    model=model,
    max_steps=3
)

GradioUI(agent).launch()