from smolagents.tools import Tool
from pytube import YouTube
import os

class YouTubeAudioDownloadTool(Tool):
    name = "youtube_audio_download"
    description = "Downloads the audio from a YouTube video and returns the local file path."
    inputs = {"youtube_url": {"type": "string", "description": "URL of the YouTube video"}}
    output_type = "string"
    
    def __init__(self, download_dir: str = "./downloads"):
        super().__init__()
        self.download_dir = download_dir
        os.makedirs(self.download_dir, exist_ok=True)

    def forward(self, youtube_url: str) -> str:
        yt = YouTube(youtube_url)
        stream = yt.streams.filter(only_audio=True).first()
        output_path = stream.download(output_path=self.download_dir, filename="temp_audio.mp3")
        return output_path