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

# Instantiate the DuckDuckGo search tool provided by smolagents.
duckduckgo_tool = DuckDuckGoSearchTool()

# Final answer tool (must be included)
final_answer = FinalAnswerTool()

# Define the model configuration.
model = HfApiModel(
    max_tokens=1000,  # Reduced max_tokens to help prevent token limit issues.
    temperature=0.5,
    model_id='Qwen/Qwen2.5-Coder-32B-Instruct',
    custom_role_conversions=None,
)

# Optionally, load an additional tool from the Hub (e.g., a text-to-image tool)
image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)

# Load prompt templates from YAML.
with open("prompts.yaml", 'r') as stream:
    prompt_templates = yaml.safe_load(stream)
    
# Initialize the CodeAgent with the DuckDuckGo search tool.
agent = CodeAgent(
    model=model,
    tools=[final_answer, duckduckgo_tool],
    max_steps=3,  # Reduced steps to limit the chain-of-thought length.
    verbosity_level=1,
    grammar=None,
    planning_interval=None,
    name="SearchAgent",
    description="An AI agent that uses DuckDuckGo search to fetch information from the web.",
    prompt_templates=prompt_templates
)

# Launch the Gradio UI to interact with the agent.
GradioUI(agent).launch()