from smolagents import CodeAgent,HfApiModel,load_tool,tool
import datetime
import requests
import pytz
import yaml
from tools.final_answer import FinalAnswerTool



@tool
def add_numbers(a: int, b: int) -> int:
    """
    Add two numbers together.

    Args:
        a (int): The first number.
        b (int): The second number.

    Returns:
        int: The sum of the two numbers.
    """
    return a + b
       
@tool
def get_current_time_in_timezone(timezone: str) -> str:
    """
    Get the current time in a given timezone.

    Args:
        timezone (str): A valid timezone string (e.g., 'America/New_York').

    Returns:
        str: The current time in that timezone, formatted as a string.
    """
    from datetime import datetime
    import pytz

    tz = pytz.timezone(timezone)
    now = datetime.now(tz)
    return now.strftime("%Y-%m-%d %H:%M:%S %Z")

# If the agent does not answer, the model is overloaded, please use another model or the following Hugging Face Endpoint that also contains qwen2.5 coder:
# model_id='https://pflgm2locj2t89co.us-east-1.aws.endpoints.huggingface.cloud' 

model = HfApiModel(
    model_id='Qwen/Qwen2.5-Coder-32B-Instruct',# it is possible that this model may be overloaded
    max_tokens=1024,
    temperature=0.5
)

with open("prompts.yaml", 'r') as stream:
    prompt_templates = yaml.safe_load(stream)

tool=[
    "add_numbers",
    "get_current_time_in_timezone", "FinalAnswerTool"
]

agent = CodeAgent(
    model=model,
    tools=[add_numbers, get_current_time_in_timezone],
    max_steps=6,
    verbosity_level=1,
    name="PracticeAgent",
    description="Agent with simple math and time tools",
    prompt_templates=prompt_templates
)


# Entry Point


if __name__ == "__main__":
   print(agent.run("What is 6 + 99, and also tell me the current time in US/Pacific"))