from smolagents import CodeAgent, HfApiModel
from tools.mood_to_need import mood_to_need
from tools.need_to_destination import need_to_destination
from tools.get_weather import get_weather
from tools.get_flights import get_flights
from tools.final_wrap import final_wrap
from tools.final_answer import FinalAnswerTool
from Gradio_UI import GradioUI
import yaml

# Initialize Claude model via Hugging Face
model = HfApiModel(
    model_id="claude-3-opus-20240229",
    temperature=0.7,
    max_tokens=2048
)

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

# Define the agent with all tools
agent = CodeAgent(
    model=model,
    tools=[
        mood_to_need,          # Step 1: Mood → Need
        need_to_destination,   # Step 2: Need → Destination
        get_weather,           # Step 3: Weather for destination
        get_flights,           # Step 4: Destination → Flights
        final_wrap,            # Step 5: Claude wrap
        FinalAnswerTool()      # Required final output
    ],
    max_steps=6,
    verbosity_level=1,
    prompt_templates=prompt_templates
)

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