# app.py
from smolagents import CodeAgent, HfApiModel
from tools.final_answer import FinalAnswerTool
import yaml
from Gradio_UI import GradioUI

def create_agent():
    # Initialize the final answer tool
    final_answer = FinalAnswerTool()
    # Load prompt templates
    with open("prompts.yaml", 'r') as stream:
        prompt_templates = yaml.safe_load(stream)
    
    # Create the model
    model = HfApiModel(
        model_id='Qwen/Qwen2.5-Coder-32B-Instruct',
        max_tokens=2096,
        temperature=0.5
    )
    
    # Create and return the agent
    return CodeAgent(
        model=model,
        tools=[final_answer],
        max_steps=6,
        verbosity_level=1,
        prompt_templates=prompt_templates
    )

if __name__ == "__main__":
    # Create and launch the agent
    agent = create_agent()
    GradioUI(agent).launch(
        server_name="0.0.0.0",
        server_port=7860
    )