from smolagents import CodeAgent, HfApiModel
from tools.mood_to_need import MoodToNeedTool
from tools.need_to_destination import NeedToDestinationTool
from tools.weather_tool import WeatherTool
from tools.find_flight import FlightsFinderTool
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=[
        MoodToNeedTool(),          # Step 1: Mood → Need
        NeedToDestinationTool(),   # Step 2: Need → Destination
        WeatherTool(),             # Step 3: Weather for destination
        FlightsFinderTool(),       # Step 4: Destination → Flights           # 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()