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

# Below is an example of a tool that does nothing. Amaze us with your creativity !
@tool
def my_custom_tool(arg1:str, arg2:int)-> str: #it's import to specify the return type
    #Keep this format for the description / args / args description but feel free to modify the tool
    """A tool that does nothing yet 
    Args:
        arg1: the first argument
        arg2: the second argument
    """
    return "What magic will you build ?"



@tool
def calculator(number1: float, number2: float, operation: str) -> float:
    """对两个数字进行基础算术运算.

    Args:
        number1: 第一个数字。
        number2: 第二个数字。
        operation: 运算类型，可填写 add、subtract、multiply 或 divide
    """
    if operation == "add":
        return number1 + number2

    elif operation == "subtract":
        return number1 - number2

    elif operation == "multiply":
        return number1 * number2

    elif operation == "divide":
        if number2 == 0:
            raise ValueError("Cannot divide by zero.")
        return number1 / number2

    else:
        raise ValueError(
            "Unsupported operation. Use add, subtract, multiply, or divide."
        )

    

@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 in that timezone
        local_time = datetime.datetime.now(tz).strftime("%Y-%m-%d %H:%M:%S")
        return f"The current local time in {timezone} is: {local_time}"
    except Exception as e:
        return f"Error fetching time for timezone '{timezone}': {str(e)}"



# 创建一个底层网页搜索工具
web_search = DuckDuckGoSearchTool(
    max_results=8
)


@tool
def search_today_news(topic: str) -> str:
    """搜索指定主题的今日新闻，供 Agent 整理成1至2段新闻简报。

    Args:
        topic: 新闻主题，例如“人工智能”“全球财经”“中国科技”。
            输入“综合”时，搜索国内外综合新闻。
    """
    # 获取北京时间的当前日期
    timezone = pytz.timezone("Asia/Shanghai")
    today = datetime.datetime.now(timezone).strftime("%Y-%m-%d")

    # 根据主题组织搜索语句
    if topic.strip() == "综合":
        query = (
            f"{today} 今日重要新闻 "
            f"国内新闻 国际新闻 科技新闻 财经新闻"
        )
    else:
        query = f"{today} {topic} 今日最新新闻"

    # 调用 DuckDuckGo 搜索
    search_results = web_search(query)

    # 将搜索结果返回给 Agent
    return f"""
当前日期：{today}
新闻主题：{topic}

以下是互联网搜索结果：

{search_results}

请根据以上搜索结果，筛选最重要的新闻，
用中文写成1至2段简短新闻简报。
不要逐条罗列，不要编造搜索结果中没有的信息。
"""






final_answer = FinalAnswerTool()

# 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(
max_tokens=2096,
temperature=0.5,
model_id='Qwen/Qwen2.5-Coder-32B-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)

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


GradioUI(agent).launch()