from smolagents import CodeAgent, DuckDuckGoSearchTool, HfApiModel, load_tool, tool
import yaml
from tools.final_answer import FinalAnswerTool
from Gradio_UI import GradioUI

def get_ai_agent_info(tool: DuckDuckGoSearchTool) -> str:
    """Searches the web for information about becoming an AI agent."""
    search_query = "becoming an AI agent"
    return tool.search(search_query)

def get_ai_agent_roadmap(model: HfApiModel) -> str:
    """Generates a roadmap for becoming an AI agent using AI model."""
    prompt = "Provide a detailed roadmap for becoming an AI agent."
    return model.generate(prompt)

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

# Load tools
final_answer = FinalAnswerTool()
search_tool = DuckDuckGoSearchTool()
image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)

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

# Create AI Agent
agent = CodeAgent(
    model=model,
    tools=[final_answer, search_tool, image_generation_tool],
    max_steps=6,
    verbosity_level=1,
    name="AI Career Guide",
    description="An AI-powered assistant for learning how to become an AI agent.",
    prompt_templates=prompt_templates
)

# Launch UI
GradioUI(agent).launch()