Text-to-Document-Generation/PDF-Redaction-API

1"""2FastAPI application for PDF redaction using NER3"""4from fastapi import FastAPI, File, UploadFile, HTTPException, BackgroundTasks5from fastapi.responses import FileResponse6from fastapi.middleware.cors import CORSMiddleware7from pydantic import BaseModel8from typing import List, Optional, Dict9import uvicorn10import os11import uuid12import shutil13from pathlib import Path14import logging15import sys16from app.redaction import PDFRedactor17from client_supabase import supabase  # Supabase client in separate file18 19# Configure logging20logging.basicConfig(21    level=logging.INFO,22    stream=sys.stdout,23    force=True,24    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"25)26logger = logging.getLogger(__name__)27 28# Initialize FastAPI app29app = FastAPI(30    title="PDF Redaction API",31    description="Redact sensitive information from PDFs using Named Entity Recognition",32    version="1.0.0"33)34 35# CORS middleware36app.add_middleware(37    CORSMiddleware,38    allow_origins=["*"],39    allow_credentials=True,40    allow_methods=["*"],41    allow_headers=["*"],42)43 44# Create directories45UPLOAD_DIR = Path("uploads")46OUTPUT_DIR = Path("outputs")47UPLOAD_DIR.mkdir(exist_ok=True)48OUTPUT_DIR.mkdir(exist_ok=True)49 50# Initialize redactor51redactor = PDFRedactor()52 53# ---------------- Response Models ----------------54class RedactionEntity(BaseModel):55    entity_type: str56    entity_text: str57    page: int58    word_count: int59 60class RedactionResponse(BaseModel):61    job_id: str62    status: str63    message: str64    entities: Optional[List[RedactionEntity]] = None65    redacted_file_url: Optional[str] = None66 67class RedactionStatusResponse(BaseModel):68    request_id: str69    status: str70    files: List[str]71    message: str72 73class HealthResponse(BaseModel):74    status: str75    version: str76    model_loaded: bool77 78# ---------------- DB Status Helpers ----------------79def set_request_status(request_id: str, status: str):80    """Update the status column in document_requests for the given request_id."""81    supabase.from_("document_requests").update({"status": status}).eq("id", request_id).execute()82    logger.info(f"Request {request_id} status -> {status}")83 84def get_request_status(request_id: str) -> str:85    """Fetch current status from document_requests."""86    response = (87        supabase88        .from_("document_requests")89        .select("status")90        .eq("id", request_id)91        .maybe_single()92        .execute()93    )94    if response.data:95        return response.data["status"]96    return "not_found"97 98# ---------------- Helper Functions ----------------99def get_public_url(bucket: str, storage_path: str) -> str:100    return f"{os.getenv('SUPABASE_URL')}/storage/v1/object/public/{bucket}/{storage_path}"101 102def cleanup_files(job_id: str):103    """Clean up temporary files after a delay"""104    try:105        upload_path = UPLOAD_DIR / f"{job_id}.pdf"106        if upload_path.exists():107            upload_path.unlink()108        logger.info(f"Cleaned up files for job {job_id}")109    except Exception as e:110        logger.error(f"Error cleaning up files for job {job_id}: {str(e)}")111 112def cleanup_temp_files(paths: List[Path]):113    for path in paths:114        if path.exists():115            path.unlink()116 117def download_file_from_supabase(bucket: str, storage_path: str, local_path: Path):118    logger.info(f"Downloading {storage_path} to {local_path}")119    data = supabase.storage.from_(bucket).download(storage_path)120    if not data:121        raise Exception(f"Failed to download {storage_path}")122    with local_path.open("wb") as f:123        f.write(data)124 125def upload_file_to_supabase(bucket: str, storage_path: str, local_path: Path):126    logger.info(f"Uploading {local_path} to {storage_path}")127    with local_path.open("rb") as f:128        content = f.read()129    supabase.storage.from_(bucket).upload(130        path=storage_path,131        file=content,132        file_options={133            "upsert": "true",134            "content-type": "application/pdf"135        }136    )137 138def redact_request(request_id: str, bucket: str = "doc_storage"):139    """140    Background task: redact all files for a given request_id.141    DB writes: 2 total — one at start (redacting), one at end (redacted | failed).142    The 'pending' write is done by the endpoint before this task is dispatched.143    """144    try:145        print("Request arrived at redact_request function")146        # Write 1: mark as redacting147        set_request_status(request_id, "redacting")148 149        response = (150            supabase151            .from_("request_files")152            .select("id, storage_path")153            .eq("request_id", request_id)154            .eq("file_role","seed")155            .execute()156        )157 158        files = response.data159        if not files:160            set_request_status(request_id, "approved")161            raise Exception(f"No files found for request {request_id}")162 163        for file in files:164            storage_path = file["storage_path"]165            local_upload = UPLOAD_DIR / f"{uuid.uuid4()}.pdf"166            local_output = OUTPUT_DIR / f"{uuid.uuid4()}_redacted.pdf"167 168            download_file_from_supabase(bucket, storage_path, local_upload)169            redactor.redact_document(pdf_path=str(local_upload), output_path=str(local_output))170            upload_file_to_supabase(bucket, storage_path, local_output)171            cleanup_temp_files([local_upload, local_output])172 173        # Write 2: mark as redacted174        set_request_status(request_id, "redacted")175 176    except Exception as e:177        print(f"Redaction failed for {request_id}: {str(e)}")178        logger.error(f"Redaction failed for {request_id}: {str(e)}")179        # Write 2 (error path): mark as failed180        set_request_status(request_id, "failed")181 182# ----------------- Existing Endpoints -----------------183@app.get("/", response_model=HealthResponse)184async def root():185    return HealthResponse(186        status="healthy",187        version="1.0.0",188        model_loaded=redactor.is_model_loaded()189    )190 191@app.get("/health", response_model=HealthResponse)192async def health_check():193    return HealthResponse(194        status="healthy",195        version="1.0.0",196        model_loaded=redactor.is_model_loaded()197    )198 199@app.post("/redact", response_model=RedactionResponse)200async def redact_pdf(201    background_tasks: BackgroundTasks,202    file: UploadFile = File(...),203    dpi: int = 300,204    entity_types: Optional[str] = None205):206    if not file.filename.endswith('.pdf'):207        raise HTTPException(status_code=400, detail="Only PDF files are supported")208    job_id = str(uuid.uuid4())209    upload_path = UPLOAD_DIR / f"{job_id}.pdf"210    output_path = OUTPUT_DIR / f"{job_id}_redacted.pdf"211    try:212        with upload_path.open("wb") as buffer:213            shutil.copyfileobj(file.file, buffer)214 215        entity_filter = None216        if entity_types:217            entity_filter = [et.strip() for et in entity_types.split(',')]218 219        result = redactor.redact_document(220            pdf_path=str(upload_path),221            output_path=str(output_path),222            dpi=dpi,223            entity_filter=entity_filter224        )225 226        response_entities = [227            RedactionEntity(228                entity_type=e['entity_type'],229                entity_text=e['entity_text'],230                page=e['words'][0]['page'] if e['words'] else 0,231                word_count=len(e['words'])232            ) for e in result['entities']233        ]234 235        background_tasks.add_task(cleanup_files, job_id)236 237        return RedactionResponse(238            job_id=job_id,239            status="completed",240            message=f"Successfully redacted {len(result['entities'])} entities",241            entities=response_entities,242            redacted_file_url=f"/download/{job_id}"243        )244 245    except Exception as e:246        logger.error(f"Error processing job {job_id}: {str(e)}")247        if upload_path.exists():248            upload_path.unlink()249        if output_path.exists():250            output_path.unlink()251        raise HTTPException(status_code=500, detail=f"Error processing PDF: {str(e)}")252 253@app.get("/download/{job_id}")254async def download_redacted_pdf(job_id: str):255    output_path = OUTPUT_DIR / f"{job_id}_redacted.pdf"256    if not output_path.exists():257        raise HTTPException(status_code=404, detail="Redacted file not found")258    return FileResponse(259        path=output_path,260        media_type="application/pdf",261        filename=f"redacted_{job_id}.pdf"262    )263 264@app.delete("/cleanup/{job_id}")265async def cleanup_job(job_id: str):266    try:267        cleanup_files(job_id)268        output_path = OUTPUT_DIR / f"{job_id}_redacted.pdf"269        if output_path.exists():270            output_path.unlink()271        return {"message": f"Successfully cleaned up files for job {job_id}"}272    except Exception as e:273        raise HTTPException(status_code=500, detail=f"Error cleaning up: {str(e)}")274 275@app.get("/stats")276async def get_stats():277    upload_count = len(list(UPLOAD_DIR.glob("*.pdf")))278    output_count = len(list(OUTPUT_DIR.glob("*.pdf")))279    return {280        "pending_uploads": upload_count,281        "processed_files": output_count,282        "model_loaded": redactor.is_model_loaded()283    }284 285# ----------------- NEW Endpoints -----------------286@app.post("/redact_by_request/{request_id}", response_model=RedactionStatusResponse)287async def redact_by_request(request_id: str, background_tasks: BackgroundTasks):288    # Check current DB status to avoid re-triggering an in-progress job289    current_status = get_request_status(request_id)290 291    if current_status == "redacting":292        return RedactionStatusResponse(293            request_id=request_id,294            status="redacting",295            files=[],296            message="Redaction already in progress"297        )298 299    # Write 1: set pending before dispatching background task300    set_request_status(request_id, "pending")301    background_tasks.add_task(redact_request, request_id)302 303    return RedactionStatusResponse(304        request_id=request_id,305        status="pending",306        files=[],307        message="Redaction started in background"308    )309 310@app.get("/redaction_status/{request_id}", response_model=RedactionStatusResponse)311async def get_redaction_status(request_id: str):312    status = get_request_status(request_id)313 314    files: List[str] = []315 316    if status == "redacted":317        response = (318            supabase319            .from_("request_files")320            .select("storage_path")321            .eq("file_role","seed")322            .eq("request_id", request_id)323            .execute()324        )325        if response.data:326            files = [327                get_public_url("doc_storage", row["storage_path"])328                for row in response.data329            ]330 331    message = {332        "redacted": "Redaction completed",333        "pending": "Redaction pending",334        "redacting": "Redaction in progress",335        "failed": "Redaction failed",336        "not_found": "Request not found",337    }.get(status, status)338 339    return RedactionStatusResponse(340        request_id=request_id,341        status=status,342        files=files,343        message=message344    )