Update app.py

app.py

@@ -5,35 +5,38 @@ import uuid
 import logging
 import asyncio
 from pathlib import Path
+from typing import Dict
 
 import httpx
-from fastapi import FastAPI, HTTPException, Request, BackgroundTasks
+from fastapi import FastAPI, HTTPException, Request, BackgroundTasks, status
 from fastapi.responses import JSONResponse
 from fastapi.staticfiles import StaticFiles
 from yt_dlp import YoutubeDL
 
 # --- Basic Configuration ---
-logging.basicConfig(
-    level=logging.INFO,
-    format="%(asctime)s - %(levelname)s - %(message)s",
-    handlers=[logging.StreamHandler()]
-)
+logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
 BASE_URL = os.getenv("BASE_URL")
 TEMP_DIR = Path("/tmp/downloads")
 STATIC_DIR = Path("static")
 TEMP_DIR.mkdir(exist_ok=True)
 STATIC_DIR.mkdir(exist_ok=True)
-FILE_LIFETIME_SECONDS = 900
+FILE_LIFETIME_SECONDS = 1800  # 30 minutes
+
+# --- In-memory store for task statuses ---
+# In a real-world, scalable application, you'd use a database like Redis or a message queue.
+# For a Hugging Face Space, this simple dictionary is sufficient.
+task_statuses: Dict[str, Dict] = {}
+
 
 # --- FastAPI App Initialization ---
 app = FastAPI(
-    title="Video Processing API v3",
-    description="Generates a temporary download link for a merged YouTube video, optimized for speed."
+    title="Async Video Processor",
+    description="An API to process videos asynchronously without timeouts."
 )
 app.mount("/static", StaticFiles(directory=STATIC_DIR), name="static")
 
 
-# --- Helper Functions ---
+# --- Helper Functions and Background Worker ---
 async def cleanup_file(filepath: Path):
     await asyncio.sleep(FILE_LIFETIME_SECONDS)
     try:
@@ -44,16 +47,14 @@ async def cleanup_file(filepath: Path):
         logging.error(f"Error during cleanup of {filepath.parent}: {e}")
 
 def get_best_formats_with_fallback(data: dict, requested_quality: int):
-    if "formats" not in data:
-        raise ValueError("The 'formats' key is missing from the Info API response.")
+    # This function remains the same as before.
+    if "formats" not in data: raise ValueError("The 'formats' key is missing from the Info API response.")
     video_url, audio_url = None, None
     video_formats = [f for f in data["formats"] if f.get("vcodec") != "none" and f.get("acodec") == "none" and f.get("ext") == "mp4" and f.get("height")]
     video_formats.sort(key=lambda f: f["height"], reverse=True)
     for f in video_formats:
-        if f["height"] <= requested_quality:
-            video_url = f["url"]; logging.info(f"Selected video quality: {f['height']}p (requested <= {requested_quality}p)"); break
-    if not video_url and video_formats:
-        f = video_formats[-1]; video_url = f["url"]; logging.warning(f"Requested quality not available. Falling back to lowest available: {f['height']}p")
+        if f["height"] <= requested_quality: video_url = f["url"]; logging.info(f"Selected video quality: {f['height']}p (requested <= {requested_quality}p)"); break
+    if not video_url and video_formats: f = video_formats[-1]; video_url = f["url"]; logging.warning(f"Requested quality not available. Falling back to lowest available: {f['height']}p")
     audio_formats = [f for f in data.get("formats", []) if f.get("acodec") not in (None, "none") and f.get("vcodec") == "none"]
     if audio_formats:
         audio_formats.sort(key=lambda f: f.get("abr", 0) or 0, reverse=True)
@@ -65,91 +66,100 @@ def get_best_formats_with_fallback(data: dict, requested_quality: int):
         raise ValueError("Could not find suitable video and/or audio streams from the Info API.")
     return video_url, audio_url
 
-# This is a synchronous helper function for downloading
 def run_ytdlp_download(url: str, out_path: Path):
-    """
-    Runs a single yt-dlp download. This is a blocking operation.
-    """
-    ydl_opts = {
-        'outtmpl': str(out_path),
-        # OPTIMIZATION: These options silence the percentage-based download logs
-        'quiet': True,
-        'noprogress': True,
-    }
+    ydl_opts = {'outtmpl': str(out_path), 'quiet': True, 'noprogress': True}
     with YoutubeDL(ydl_opts) as ydl:
         ydl.download([url])
 
-# --- API Endpoints ---
-@app.get("/")
-def read_root():
-    return {"message": "Video Processing API v3 is running."}
-
-@app.post("/api/process")
-async def process_video(request: Request, background_tasks: BackgroundTasks):
-    if not BASE_URL:
-        logging.error("FATAL: BASE_URL is not configured in the server environment.")
-        raise HTTPException(status_code=500, detail="Server is not configured correctly.")
-    body = await request.json()
-    video_url = body.get("url")
-    try:
-        quality = int(body.get("quality", "1080"))
-    except (ValueError, TypeError):
-        raise HTTPException(status_code=400, detail="'quality' must be a valid number (e.g., 1080, 720).")
-    if not video_url:
-        raise HTTPException(status_code=400, detail="A 'url' is required.")
-    logging.info(f"Received request for URL: {video_url} with quality: {quality}p")
+async def process_in_background(
+    task_id: str,
+    video_url: str,
+    quality: int,
+    base_url_for_links: str,
+    background_tasks: BackgroundTasks
+):
+    """
+    This is the main worker function that runs in the background.
+    """
+    task_statuses[task_id] = {"status": "processing", "message": "Fetching video info..."}
     try:
+        # Step 1: Get video info
         async with httpx.AsyncClient(follow_redirects=True) as client:
             info_api_url = f"{BASE_URL.rstrip('/')}/api/info"
             response = await client.get(info_api_url, params={"url": video_url, "playlist": "false"}, timeout=30.0)
             response.raise_for_status()
             video_data = response.json()
             video_stream_url, audio_stream_url = get_best_formats_with_fallback(video_data, quality)
-    except Exception as e:
-        logging.error(f"Error processing video info: {e}")
-        raise HTTPException(status_code=500, detail=f"An unexpected error occurred: {e}")
-
-    task_id = str(uuid.uuid4())
-    final_output_dir = STATIC_DIR / task_id
-    final_output_dir.mkdir()
-    final_output_path = final_output_dir / "video.mp4"
-    video_path = TEMP_DIR / f"{task_id}_video.mp4"
-    audio_path = TEMP_DIR / f"{task_id}_audio.m4a"
 
-    try:
-        # --- SPEED OPTIMIZATION ---
-        # Run the two downloads concurrently instead of one after the other.
-        logging.info(f"Starting concurrent download for task {task_id}")
-
-        # Create two tasks to run our synchronous download function in separate threads
-        video_download_task = asyncio.to_thread(run_ytdlp_download, video_stream_url, video_path)
-        audio_download_task = asyncio.to_thread(run_ytdlp_download, audio_stream_url, audio_path)
+        # Step 2: Download files
+        task_statuses[task_id]["message"] = "Downloading video and audio streams..."
+        final_output_dir = STATIC_DIR / task_id
+        final_output_dir.mkdir()
+        video_path = TEMP_DIR / f"{task_id}_video.mp4"
+        audio_path = TEMP_DIR / f"{task_id}_audio.m4a"
 
-        # Wait for both downloads to complete
-        await asyncio.gather(video_download_task, audio_download_task)
+        video_dl_task = asyncio.to_thread(run_ytdlp_download, video_stream_url, video_path)
+        audio_dl_task = asyncio.to_thread(run_ytdlp_download, audio_stream_url, audio_path)
+        await asyncio.gather(video_dl_task, audio_dl_task)
 
-        logging.info(f"Download complete. Starting FFmpeg merge for task {task_id}")
+        # Step 3: Merge files
+        task_statuses[task_id]["message"] = "Merging files with FFmpeg..."
+        final_output_path = final_output_dir / "video.mp4"
         subprocess.run(
             ['ffmpeg', '-i', str(video_path), '-i', str(audio_path), '-c', 'copy', str(final_output_path)],
             check=True, capture_output=True, text=True
         )
-        logging.info(f"Merge successful. Final file at: {final_output_path}")
 
-    except Exception as e:
-        logging.error(f"Download or Merge Failed for task {task_id}: {e}")
-        shutil.rmtree(final_output_dir)
-        raise HTTPException(status_code=500, detail=f"Failed during file processing: {e}")
-    finally:
-        if video_path.exists(): video_path.unlink()
-        if audio_path.exists(): audio_path.unlink()
-
-    download_url = request.url_for('static', path=f"{task_id}/video.mp4")
-    background_tasks.add_task(cleanup_file, final_output_path)
-    logging.info(f"Responding with download URL: {download_url}")
-    return JSONResponse(
-        content={
-            "success": True,
-            "download_url": str(download_url),
+        # Step 4: Finalize and set status to complete
+        download_url = f"{base_url_for_links.rstrip('/')}/static/{task_id}/video.mp4"
+        task_statuses[task_id] = {
+            "status": "complete",
+            "download_url": download_url,
             "expires_in": f"{FILE_LIFETIME_SECONDS} seconds"
         }
-    )
+        background_tasks.add_task(cleanup_file, final_output_path)
+
+    except Exception as e:
+        logging.error(f"Task {task_id} failed: {e}")
+        task_statuses[task_id] = {"status": "failed", "error": str(e)}
+    finally:
+        if 'video_path' in locals() and video_path.exists(): video_path.unlink()
+        if 'audio_path' in locals() and audio_path.exists(): audio_path.unlink()
+
+
+# --- API Endpoints ---
+@app.post("/api/start-processing", status_code=status.HTTP_202_ACCEPTED)
+async def start_processing_job(request: Request, background_tasks: BackgroundTasks):
+    """
+    Accepts a job and starts it in the background. Returns a task ID immediately.
+    """
+    body = await request.json()
+    video_url = body.get("url")
+    quality = int(body.get("quality", "1080"))
+
+    if not video_url:
+        raise HTTPException(status_code=400, detail="A 'url' is required.")
+
+    task_id = str(uuid.uuid4())
+    task_statuses[task_id] = {"status": "queued"}
+
+    # We need the base URL of our own app to construct the final download link
+    base_url_for_links = str(request.base_url)
+
+    background_tasks.add_task(
+        process_in_background, task_id, video_url, quality, base_url_for_links, background_tasks
+    )
+
+    status_url = request.url_for('get_job_status', task_id=task_id)
+    return {"task_id": task_id, "status_url": str(status_url)}
+
+
+@app.get("/api/status/{task_id}")
+async def get_job_status(task_id: str):
+    """
+    Allows the client to poll for the status of a background job.
+    """
+    status_info = task_statuses.get(task_id)
+    if not status_info:
+        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Task not found")
+    return status_info