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

# Initialize the mandatory final answer tool
final_answer = FinalAnswerTool()

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

# Import the image generation tool from the Hugging Face Hub
image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)

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

# Create the agent and provide it with the correct tools
agent = CodeAgent(
    model=model,
    # Ensure both the image tool and final answer tool are in this list
    tools=[image_generation_tool, final_answer],
    max_steps=6,
    verbosity_level=1,
    prompt_templates=prompt_templates
)

# Launch the user interface
GradioUI(agent).launch()