Spaces:
Sleeping
Sleeping
File size: 2,852 Bytes
0eb5f56 a798955 0eb5f56 03cf704 a798955 0eb5f56 03cf704 0eb5f56 03cf704 0eb5f56 03cf704 0eb5f56 03cf704 0eb5f56 7e5b577 8e46fec dc37e9e 7e5b577 a798955 7e5b577 dc37e9e a798955 dc37e9e a798955 dc37e9e 7e5b577 a798955 7e5b577 dc37e9e 7e5b577 0eb5f56 |
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 |
import os
import asyncio
from fastapi import FastAPI, Request, Form
from fastapi.responses import HTMLResponse, JSONResponse, FileResponse
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates
from helium import start_chrome, kill_browser, go_to
import uuid
import logging
import pathlib
import time
# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
# Create FastAPI app
app = FastAPI()
# Determine base directory
BASE_DIR = pathlib.Path(__file__).parent.parent
# Set up templates and static files
templates_path = BASE_DIR / "templates"
static_path = BASE_DIR / "static"
screenshots_path = BASE_DIR / "screenshots"
templates = Jinja2Templates(directory=str(templates_path))
app.mount("/static", StaticFiles(directory=str(static_path)), name="static")
# Create screenshots directory if it doesn't exist
os.makedirs(str(screenshots_path), exist_ok=True)
# Mount the screenshots directory
app.mount("/screenshots", StaticFiles(directory=str(screenshots_path)), name="screenshots")
@app.get("/", response_class=HTMLResponse)
async def read_root(request: Request):
return templates.TemplateResponse("index.html", {"request": request})
@app.post("/take-screenshot")
async def take_screenshot(url: str = Form(...)):
logger.info(f"Taking screenshot of URL: {url}")
try:
# Generate a unique filename
filename = f"{uuid.uuid4()}.png"
filepath = str(screenshots_path / filename)
# Log environment info
logger.info(f"HOME env: {os.environ.get('HOME')}")
logger.info(f"Current working directory: {os.getcwd()}")
# Launch Chrome with simple options
logger.info("Launching Chrome browser")
driver = start_chrome(headless=True)
try:
# Navigate to URL
logger.info(f"Navigating to URL: {url}")
go_to(url)
# Wait for page to load
logger.info("Waiting for page to load")
time.sleep(3)
# Take the screenshot
logger.info(f"Taking screenshot at {filepath}")
driver.save_screenshot(filepath)
logger.info(f"Screenshot saved to {filepath}")
finally:
# Always close the browser
kill_browser()
logger.info("Browser closed")
return JSONResponse({
"success": True,
"screenshot_url": f"/screenshots/{filename}"
})
except Exception as e:
logger.error(f"Error taking screenshot: {str(e)}")
return JSONResponse({
"success": False,
"error": str(e)
}, status_code=500)
# Add a health check endpoint
@app.get("/health")
async def health_check():
return {"status": "ok"} |