Spaces:
Sleeping
Sleeping
File size: 5,324 Bytes
34d0dc9 a1a9708 34d0dc9 a089552 34d0dc9 a089552 34d0dc9 a089552 34d0dc9 a089552 34d0dc9 d57e0f9 34d0dc9 a089552 34d0dc9 a089552 34d0dc9 a089552 34d0dc9 a089552 34d0dc9 a089552 34d0dc9 a089552 34d0dc9 a1a9708 34d0dc9 a089552 34d0dc9 a089552 34d0dc9 a1a9708 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 | import os
import asyncio
import uuid
import random
from app.scraper_manager import ScraperManager
from fastapi import APIRouter, BackgroundTasks, HTTPException
from fastapi.responses import FileResponse, Response
from pathlib import Path
from pydantic import BaseModel
from typing import List
from loguru import logger
router = APIRouter()
log = logger.bind(context="API")
jobs = {}
download_lock = asyncio.Lock()
class ScrapeRequest(BaseModel):
urls: List[str]
async def process_url(
job_id: str,
url_id: str,
url: str,
delay: float = 0.0,
limit: asyncio.Semaphore = asyncio.Semaphore(10),
):
async with limit:
jobs[job_id]["urls"][url_id]["status"] = "processing"
if delay > 0:
await asyncio.sleep(delay)
manager = ScraperManager()
try:
result = await manager.scrape(url)
jobs[job_id]["urls"][url_id]["status"] = "completed"
return result
except Exception as e:
log.exception(f"Error scraping {url}: {e}")
if "Bot Blocked" in str(e):
jobs[job_id]["urls"][url_id]["status"] = "Bot Blocked"
elif "Wrong URL" in str(e):
jobs[job_id]["urls"][url_id]["status"] = "Wrong URL"
elif "LLM" in str(e):
jobs[job_id]["urls"][url_id]["status"] = "LLM limit exceeded"
else:
jobs[job_id]["urls"][url_id]["status"] = "Failed"
return {"error": jobs[job_id]["urls"][url_id]["status"], "url": url}
finally:
await manager.close()
async def run(job_id: str, urls: List[str], url_tracking: dict):
jobs[job_id]["status"] = "processing"
results = []
limit = asyncio.Semaphore(10)
random_delay = random.uniform(0.0, 2.0)
# Create tasks as before...
tasks = [
process_url(job_id=job_id, url_id=url_id, url=info["url"], delay=random_delay, limit=limit)
for url_id, info in url_tracking.items()
]
# Use as_completed to update progress one by one
for task in asyncio.as_completed(tasks):
result = await task
results.append(result)
jobs[job_id]["status"] = "completed"
# Wait until the end to store the heavy result data
jobs[job_id]["data"] = results
@router.post("/scrape")
async def trigger_scrape(request: ScrapeRequest, background_tasks: BackgroundTasks):
"""
Accepts a list of URLs, generates a job ID, and starts scraping in the background.
"""
if not request.urls:
log.warning("Received scrape request with no URLs")
raise HTTPException(status_code=400, detail="No URLs provided")
log.success(f"Received scrape request for {len(request.urls)} URLs")
job_id = str(uuid.uuid4())
# Store tracking by index so duplicate URLs don't overwrite each other
url_tracking = {str(i): {"url": url, "status": "pending"} for i, url in enumerate(request.urls)}
jobs[job_id] = {
"status": "pending",
"urls": url_tracking,
"data": None,
"error": None,
}
# Send the scraping task to the background
background_tasks.add_task(run, job_id, request.urls, url_tracking)
return {
"job_id": job_id,
"status": "pending",
"message": f"Started scraping {len(request.urls)} URLs in the background.",
}
@router.get("/results")
async def get_all_results():
"""
Return a summary of ALL jobs for the frontend to poll once
"""
return {
job_id: {
"status": info["status"],
"urls": info.get("urls", {}),
"error": info["error"],
}
for job_id, info in jobs.items()
}
DATA_DIR = Path(__file__).resolve().parent
@router.get("/download")
async def download_jsonl():
"""
Returns the data.jsonl file as a downloadable attachment and deletes it after.
Handles double-clicks gracefully by returning 204 if state was just cleared.
"""
async with download_lock:
file_path = DATA_DIR / "data.jsonl"
if not file_path.exists():
# If file is gone but jobs are cleared, it's likely a double-click
if not jobs:
log.info("Download requested but data already cleared. Returning 204.")
return Response(status_code=204)
raise HTTPException(
status_code=404, detail="No data file found. Run a scrape job first."
)
log.success("Downloading data.jsonl")
# Read the file into memory first to avoid race conditions with FileResponse
content = file_path.read_bytes()
is_any_job_active = any(
job["status"] in ["pending", "processing"] for job in jobs.values()
)
if not is_any_job_active:
file_path.unlink(missing_ok=True)
jobs.clear()
log.info("File deleted & jobs cleared: No active jobs found.")
return Response(
content=content,
media_type="application/x-ndjson",
headers={"Content-Disposition": "attachment; filename=data.jsonl"},
)
|