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

# ... (Your existing my_custom_tool and get_current_time_in_timezone definitions)

final_answer = FinalAnswerTool()
model = HfApiModel(
    max_tokens=2096,
    temperature=0.5,
    model_id='Qwen/Qwen2.5-Coder-32B-Instruct',
    custom_role_conversions=None,
)

# Initialize the DuckDuckGo Search Tool
duckduckgo_search_tool = DuckDuckGoSearchTool()

# Removed the attempt to load the Wikipedia tool
# The previous code to load the Wikipedia tool was here.

with open("prompts.yaml", 'r') as stream:
    prompt_templates = yaml.safe_load(stream)
    tools_list = [final_answer, duckduckgo_search_tool]
    # The conditional block to add the Wikipedia tool has been removed.

    agent = CodeAgent(
        model=model,
        tools=tools_list,  # Only includes final_answer and duckduckgo_search_tool now
        max_steps=6,
        verbosity_level=1,
        grammar=None,
        planning_interval=None,
        name=None,
        description=None,
        prompt_templates=prompt_templates
    )

GradioUI(agent).launch()