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

from Gradio_UI import GradioUI


@tool
def sum_tool(num1: int, num2: int) -> int:
    """ A tool that adds up the two integers
    Args:
        num1: The first integer to add
        num2: The second integer to add
    Returns:
        int: A integer contains the sum of the two integers.

    Raises:
        TypeError: If either num1 or num2 is not an integer.
    """
    if not isinstance(num1, int) or not isinstance(num2, int):
        raise TypeError("Both inputs must be integers")
        
    try:
        res = num1 + num2
        return f"The sum of two integers is: {res}"
    except Exception as e:
        return f"An error occurred: {str(e)}"


@tool
def get_current_time_in_timezone(timezone: str) -> str:
    """A tool that fetches the current local time in a specified timezone.
    Args:
        timezone: A string representing a valid timezone (e.g., 'America/New_York').
    """
    try:
        # Create timezone object
        tz = pytz.timezone(timezone)
        # Get current time rent local time in {timezone} is: {local_time}"
    except Exception as e:
        return f"Error fetching time for timezone '{timezone}': {str(e)}"


final_answer = FinalAnswerTool()
model = HfApiModel(
max_tokens=2096,
temperature=0.5,
model_id='deepseek-ai/deepseek-coder-1.3b-instruct',# it is possible that this model may be overloaded
custom_role_conversions=None,
)


# Import tool from Hub
image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)
search_tool = DuckDuckGoSearchTool()

with open("prompts.yaml", 'r') as stream:
    prompt_templates = yaml.safe_load(stream)
    
agent = CodeAgent(
    model=model,
    tools=[final_answer, search_tool, sum_tool], ## add your tools here (don't remove final answer)
    max_steps=6,
    verbosity_level=1,
    grammar=None,
    planning_interval=None,
    name=None,
    description=None,
    prompt_templates=prompt_templates
)

if agent.model.last_input_token_count is not None:
    total_input_tokens += agent.model.last_input_token_count
else:
    # Handle the case where last_input_token_count is None
    # For example, you can set total_input_tokens to a default value
    total_input_tokens = 0

GradioUI(agent).launch()