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"top {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='https://pflgm2locj2t89co.us-east-1.aws.endpoints.huggingface.cloud',  # Note: May need to switch if overloaded
    custom_role_conversions=None,
)

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

# Initialize the agent with the updated tools list and custom logic
agent = CodeAgent(
    model=model,
    tools=[construct_course_search_query, search_tool, final_answer],
    max_steps=1,  # Reduced from 9 to limit automatic progression
    planning_interval=1,  # Forces planning every step
    verbosity_level=2,  # Increased for more detailed logging
    grammar=None,
    name=None,
    description=None,
    prompt_templates=prompt_templates)

GradioUI(agent).launch()