api.py

# api.py — Document Sentinel v2.0
# Routes: POST /classify, GET /download, GET /metrics, GET /feedback, POST /evaluate
# Three-layer detection (Regex + NER + LLM) with verification and explainability

import os
import uuid
import time
from pathlib import Path

from fastapi import FastAPI, UploadFile, File, HTTPException, Query
from fastapi.responses import FileResponse, JSONResponse
from fastapi.middleware.cors import CORSMiddleware

from utils import (
    extract_text_from_pdf,
    detect_and_mask,
    redact_pdf,
    get_pipeline_metrics,
    get_feedback,
    run_evaluation,
)

app = FastAPI(
    title="Document Sentinel",
    description=(
        "Upload any PDF. "
        "Three-layer PII detection (Regex + NER + LLM) with verification loop, "
        "per-entity explainability, and performance metrics.\n\n"
        "**Layer 1 — Regex**: SSN, Email, Phone, Credit Card, IP, Dates, Financial Amounts\n\n"
        "**Layer 2 — NER (GLiNER)**: Names, Organizations, Addresses, Locations\n\n"
        "**Layer 3 — LLM**: Context-sensitive PII, NDA clauses, implicit identifiers\n\n"
        "**Verification**: Second-pass LLM confirms redaction completeness"
    ),
    version="2.0.0",
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

UPLOAD_DIR   = Path("temp_uploads")
REDACTED_DIR = Path("temp_redacted")
UPLOAD_DIR.mkdir(exist_ok=True)
REDACTED_DIR.mkdir(exist_ok=True)


# ── POST /classify ─────────────────────────────────────────────────────────────
@app.post("/classify", summary="Detect and redact PII/sensitive content in a PDF")
async def classify(file: UploadFile = File(..., description="PDF document to analyse")):
    """
    **Upload a PDF. Returns:**
    - `entities` — every sensitive entity found with:
        - `entity`: the original value
        - `label`: what type it is (EMAIL, PERSON_NAME, FINANCIAL_PENALTY, etc.)
        - `method`: how it was detected (`regex`, `ner`, or `llm`)
        - `confidence`: 0.0–1.0
        - `reason`: context-specific AI explanation of WHY this is sensitive
    - `masked_text` — full document text with all entities replaced by `[LABEL]`
    - `redacted_pdf_url` — download link for the redacted PDF
    - `layer_breakdown` — per-layer detection statistics
    - `verification` — second-pass sanitization score and any missed entities
    - `processing_time_ms` — end-to-end latency
    """
    if not file.filename.lower().endswith(".pdf"):
        raise HTTPException(status_code=400, detail="Only PDF files are accepted.")

    t_start  = time.time()
    doc_id   = str(uuid.uuid4())
    in_path  = UPLOAD_DIR   / f"{doc_id}_original.pdf"
    out_path = REDACTED_DIR / f"{doc_id}_redacted.pdf"

    # Save upload
    contents = await file.read()
    with open(in_path, "wb") as f:
        f.write(contents)

    # Step 1 — Extract text
    raw_text = extract_text_from_pdf(str(in_path))
    if not raw_text.strip():
        raise HTTPException(
            status_code=422,
            detail="Could not extract text from this PDF. It may be corrupt or fully image-based without OCR support."
        )

    # Step 2 — Detect + mask (Regex + NER + LLM + Verification)
    masked_text, entities, report = detect_and_mask(raw_text)

    # Step 3 — Redact original PDF
    redact_pdf(str(in_path), str(out_path), entities)

    elapsed_ms = round((time.time() - t_start) * 1000, 1)

    # Step 4 — Build response
    return JSONResponse({
        "document_id":        doc_id,
        "original_filename":  file.filename,
        "processing_time_ms": elapsed_ms,
        "document_type":      report.get("document_type", "unknown"),
        "sanitization_score": report.get("sanitization_score", 85.0),

        # Entity report — the "responsible AI" output
        "summary": {
            "total_entities": len(entities),
            "by_method": {
                "regex": sum(1 for e in entities if e["method"] == "regex"),
                "ner":   sum(1 for e in entities if e["method"] == "ner"),
                "llm":   sum(1 for e in entities if e["method"] == "llm"),
                "verification": sum(1 for e in entities if e["method"] == "verification"),
            },
            "by_label": _count_by_label(entities),
        },

        "layer_breakdown":  report.get("layer_breakdown", {}),
        "verification":     report.get("verification"),
        "entities":         entities,
        "masked_text":      masked_text,
        "redacted_pdf_url": f"/download/{doc_id}",
    })


# ── GET /download/{doc_id} ────────────────────────────────────────────────────
@app.get("/download/{doc_id}", summary="Download the redacted PDF")
async def download(doc_id: str):
    path = REDACTED_DIR / f"{doc_id}_redacted.pdf"
    if not path.exists():
        raise HTTPException(status_code=404, detail="Redacted PDF not found or expired.")
    return FileResponse(
        path=str(path),
        media_type="application/pdf",
        filename=f"redacted_{doc_id[:8]}.pdf",
    )


# ── GET /metrics ──────────────────────────────────────────────────────────────
@app.get("/metrics", summary="Pipeline performance metrics")
async def metrics():
    """
    Returns aggregated metrics across all processed documents:
    - Total documents processed
    - Average sanitization score
    - Layer contribution percentages (regex / ner / llm / verification)
    - Entity type distribution
    - Verification pass rate
    """
    return JSONResponse(get_pipeline_metrics())


# ── GET /feedback/{doc_id} ────────────────────────────────────────────────────
@app.get("/feedback/{doc_id}", summary="Per-entity feedback and explanations")
async def feedback(doc_id: str):
    """
    Returns per-entity feedback for a processed document:
    - Each entity with its context-specific explanation
    - Detection method and confidence
    - Verification status
    - Layer breakdown
    """
    result = get_feedback(doc_id)
    if "error" in result:
        raise HTTPException(status_code=404, detail=result["error"])
    return JSONResponse(result)


# ── POST /evaluate ────────────────────────────────────────────────────────────
@app.post("/evaluate", summary="Run evaluation on documents")
async def evaluate():
    """
    Generates synthetic documents with known PII ground truth,
    runs the full pipeline, and computes precision/recall/F1:
    - Per entity type (SSN, EMAIL, PERSON_NAME, etc.)
    - Per detection layer (regex, ner, llm)
    - Overall pipeline metrics
    """
    results = run_evaluation(n_docs=20, use_llm=True)
    return JSONResponse(results)

# ── GET /health ───────────────────────────────────────────────────────────────
@app.get("/health")
async def health():
    return {
        "status":   "ok",
        "service":  "Document Sentinel v2.0",
        "provider": os.getenv("LLM_PROVIDER", "groq"),
        "layers":   ["regex", "ner (GLiNER)", "llm", "verification"],
    }


# ── Util ──────────────────────────────────────────────────────────────────────
def _count_by_label(entities):
    counts = {}
    for e in entities:
        counts[e["label"]] = counts.get(e["label"], 0) + 1
    return counts