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

# Custom tool for constructing the search query
@tool
def construct_course_search_query(interest: str, expertise: str, budget: str) -> str:
    """Constructs a search query for finding AI courses based on user preferences.
    Args:
        interest: The area of interest within AI (e.g., 'machine learning', 'deep learning').
        expertise: The user's current level of expertise (e.g., 'beginner', 'intermediate').
        budget: The budget available for the course (e.g., '$100', 'free').
    """
    query = f"best {interest} courses for {expertise} under {budget}"
    return query

# Existing tools from the template
search_tool = DuckDuckGoSearchTool()
final_answer = FinalAnswerTool()

# Model configuration (unchanged from template)
model = HfApiModel(
    max_tokens=2096,
    temperature=0.5,
    model_id='Llama',  # Note: May need to switch if overloaded
    custom_role_conversions=None,
)

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

# Load prompt templates (unchanged from template)
with open("prompts.yaml", 'r') as stream:
    prompt_templates = yaml.safe_load(stream)

# Initialize the agent with the updated tools list
agent = CodeAgent(
    model=model,
    tools=[construct_course_search_query, search_tool, final_answer],  # Added custom tool
    max_steps=6,
    verbosity_level=1,
    grammar=None,
    planning_interval=None,
    name=None,
    description=None,
    prompt_templates=prompt_templates
)

GradioUI(agent).launch()