ytdlp-api/main.py

import os
import uuid
import uvicorn
from typing import Any, Literal, cast
from fastapi import FastAPI, HTTPException, BackgroundTasks
from fastapi.responses import FileResponse
from pydantic import BaseModel
from yt_dlp import YoutubeDL

app = FastAPI()


# Input schema because we aren't savages
class VideoRequest(BaseModel):
    url: str


class VideoDownloadRequest(VideoRequest):
    quality: Literal["best", "1080p", "720p", "480p", "360p"] = "best"


class VideoInfoResponse(BaseModel):
    id: str | None = None
    title: str | None = None
    uploader: str | None = None
    duration: int | None = None
    webpage_url: str | None = None
    thumbnail: str | None = None
    view_count: int | None = None
    upload_date: str | None = None
    extractor: str | None = None


def cleanup(path: str):
    """Deletes the file after serving because I know you won't."""
    try:
        os.remove(path)
        print(f"Deleted trash: {path}")
    except Exception as e:
        print(f"Failed to delete {path}: {e}")


def get_video_format_selector(quality: str) -> str:
    if quality == "best":
        return "bestvideo+bestaudio/best"
    height = quality.removesuffix("p")
    return f"bestvideo[height<={height}]+bestaudio/best[height<={height}]"


@app.post("/audio")
async def download_audio(req: VideoRequest, background_tasks: BackgroundTasks):
    # Random ID so concurrent requests don't overwrite each other
    file_id = str(uuid.uuid4())
    output_template = f"/tmp/{file_id}.%(ext)s"

    # yt-dlp config: best audio, force mp3 because whisper hates weird codecs
    ydl_opts: dict[str, Any] = {
        "format": "bestaudio/best",
        "outtmpl": output_template,
        "postprocessors": [
            {
                "key": "FFmpegExtractAudio",
                "preferredcodec": "mp3",
                "preferredquality": "192",
            }
        ],
        "quiet": True,
        "no_warnings": True,
    }

    try:
        with YoutubeDL(cast(Any, ydl_opts)) as ydl:
            print(f"Yoinking audio from: {req.url}")
            info = ydl.extract_info(req.url, download=True)
            # yt-dlp might change extension post-processing, so we find the actual file
            # usually it's .mp3 if we forced it, but let's be safe
            filename = ydl.prepare_filename(info).rsplit(".", 1)[0] + ".mp3"

        if not os.path.exists(filename):
            raise HTTPException(
                status_code=500, detail="Download failed, file missing. Blame YouTube."
            )

        # queue the deletion so it happens AFTER the response is sent
        background_tasks.add_task(cleanup, filename)

        return FileResponse(
            path=filename, filename=f"{file_id}.mp3", media_type="audio/mpeg"
        )

    except Exception as e:
        print(f"L: {str(e)}")
        raise HTTPException(status_code=400, detail=f"Failed to download: {str(e)}")


@app.post("/info", response_model=VideoInfoResponse)
async def get_video_info(req: VideoRequest):
    ydl_opts: dict[str, Any] = {
        "quiet": True,
        "no_warnings": True,
    }

    try:
        with YoutubeDL(cast(Any, ydl_opts)) as ydl:
            info = ydl.extract_info(req.url, download=False)

        if not info:
            raise HTTPException(status_code=404, detail="No metadata found for URL")

        return VideoInfoResponse(
            id=info.get("id"),
            title=info.get("title"),
            uploader=info.get("uploader"),
            duration=info.get("duration"),
            webpage_url=info.get("webpage_url"),
            thumbnail=info.get("thumbnail"),
            view_count=info.get("view_count"),
            upload_date=info.get("upload_date"),
            extractor=info.get("extractor"),
        )

    except HTTPException:
        raise
    except Exception as e:
        print(f"Failed metadata fetch: {str(e)}")
        raise HTTPException(
            status_code=400, detail=f"Failed to fetch metadata: {str(e)}"
        )


@app.post("/video")
async def download_video(req: VideoDownloadRequest, background_tasks: BackgroundTasks):
    file_id = str(uuid.uuid4())
    output_template = f"/tmp/{file_id}.%(ext)s"
    ydl_opts: dict[str, Any] = {
        "format": get_video_format_selector(req.quality),
        "merge_output_format": "mp4",
        "outtmpl": output_template,
        "noplaylist": True,
        "quiet": True,
        "no_warnings": True,
    }

    try:
        with YoutubeDL(cast(Any, ydl_opts)) as ydl:
            print(f"Yoinking video from: {req.url} [{req.quality}]")
            info = ydl.extract_info(req.url, download=True)
            filename = None

            requested_downloads = info.get("requested_downloads")
            if isinstance(requested_downloads, list) and requested_downloads:
                first = requested_downloads[0]
                if isinstance(first, dict):
                    filepath = first.get("filepath")
                    if isinstance(filepath, str):
                        filename = filepath

            if not filename:
                maybe_filename = info.get("_filename")
                if isinstance(maybe_filename, str):
                    filename = maybe_filename

            if not filename:
                filename = ydl.prepare_filename(info)

        if not os.path.exists(filename):
            raise HTTPException(
                status_code=500, detail="Download failed, file missing. Blame YouTube."
            )

        background_tasks.add_task(cleanup, filename)

        return FileResponse(
            path=filename, filename=f"{file_id}.mp4", media_type="video/mp4"
        )

    except HTTPException:
        raise
    except Exception as e:
        print(f"Video download failed: {str(e)}")
        raise HTTPException(
            status_code=400, detail=f"Failed to download video: {str(e)}"
        )


if __name__ == "__main__":
    # Host on 0.0.0.0 so your docker container isn't useless
    uvicorn.run(app, host="0.0.0.0", port=8000)