from smolagents import CodeAgent, HfApiModel, load_tool, tool
import yaml

from Gradio_UI import GradioUI


# -------------------- TOOLS --------------------

@tool
def final_answer(answer: str) -> str:
    """Return the final answer to the user.

    Args:
        answer: The final response to display
    """
    return answer


# ✅ Stable image tool (no manual HF token handling needed)
image_generation_tool = load_tool(
    "agents-course/text-to-image",
    trust_remote_code=True
)


# -------------------- MODEL --------------------

model = HfApiModel(
    model_id="Qwen/Qwen2.5-Coder-32B-Instruct",
    max_tokens=1024,
    temperature=0.3,
)


# -------------------- PROMPTS --------------------

with open("prompts.yaml", "r") as f:
    prompt_templates = yaml.safe_load(f)


# -------------------- AGENT --------------------

agent = CodeAgent(
    model=model,
    tools=[
        final_answer,
        image_generation_tool
    ],
    max_steps=4,
    verbosity_level=1,
    prompt_templates=prompt_templates
)


# -------------------- UI --------------------

GradioUI(agent).launch()