main.py

"""
Conversation Data Extraction System - REST API.

Endpoints:
    GET  /                    -> web dashboard
    GET  /health              -> service + model status
    POST /analyze             -> run full pipeline on a single message
    POST /predict             -> sentiment only (backward compatible)
    POST /ingest               -> start a batch ingest job, returns a job id
    GET  /ingest/status/{id}  -> progress / result of a background ingest job
    GET  /records             -> recent stored (anonymized) records
    GET  /records/export      -> stored records exported as XML, JSON or CSV
    GET  /stats               -> aggregate statistics
    GET  /models              -> default + suggested + currently loaded HF models per task

The full pipeline extracts named entities, classifies the topic, evaluates
sentiment, pseudonymizes the text and stores the structured result. Since V3,
each of the three ML steps (NER, topic classification, sentiment) can use any
compatible Hugging Face Hub model selected by the caller, instead of only the
built-in default model for that step. Since V4, stored records can be
exported in a choice of formats (XML, JSON, CSV) via `/records/export`. Since
V7, `/ingest` runs as a background job (see `jobs.py`) instead of one long
blocking request, so the dashboard can show a progress bar and elapsed timer.
"""

from __future__ import annotations

import threading
import time
from pathlib import Path
from typing import Any

from fastapi import FastAPI, Form, HTTPException, UploadFile, File
from fastapi.responses import FileResponse, Response
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel

import db
import ingest
import jobs
from pipeline import process_message, process_batch
from processors import ner, topics, sentiment, anonymizer, model_registry

app = FastAPI(
    title="Conversation Data Extraction System",
    description=(
        "Extracts named entities, topics and sentiment from chat conversations, "
        "pseudonymizes personal data (GDPR) and stores structured results."
    ),
    version="7.0.0",
)

BASE_DIR = Path(__file__).resolve().parent
STATIC_DIR = BASE_DIR / "static"
app.mount("/static", StaticFiles(directory=str(STATIC_DIR)), name="static")

# Curated lists of well-known Hugging Face models offered as suggestions in
# the UI, one per pipeline step. Any other compatible model id can still be
# supplied manually - these lists are not a whitelist.
SUGGESTED_MODELS = {
    "sentiment": [
        sentiment.DEFAULT_MODEL_NAME,
        "cardiffnlp/twitter-roberta-base-sentiment-latest",
        "distilbert-base-uncased-finetuned-sst-2-english",
        "nlptown/bert-base-multilingual-uncased-sentiment",
    ],
    "ner": [
        ner.DEFAULT_MODEL_NAME,
        "dslim/bert-large-NER",
        "Jean-Baptiste/roberta-large-ner-english",
        "dbmdz/bert-large-cased-finetuned-conll03-english",
    ],
    "topics": [
        topics.DEFAULT_MODEL_NAME,
        "MoritzLaurer/deberta-v3-base-zeroshot-v1.1-all-33",
        "valhalla/distilbart-mnli-12-3",
    ],
}


# ---------------------------------------------------------------------------
# Request models
# ---------------------------------------------------------------------------
class TextRequest(BaseModel):
    text: str
    topic_labels: list[str] | None = None
    sentiment_model: str | None = None
    ner_model: str | None = None
    topic_model: str | None = None


# ---------------------------------------------------------------------------
# Pages
# ---------------------------------------------------------------------------
@app.get("/")
def home():
    index_file = STATIC_DIR / "index.html"
    if index_file.exists():
        return FileResponse(str(index_file))
    return {"message": "Conversation Data Extraction System is running. See /docs."}


@app.get("/health")
def health():
    return {
        "status": "ok",
        "models": {
            "ner": {"name": ner.model_name(), "ready": ner.is_ready()},
            "topics": {"name": topics.model_name(), "ready": topics.is_ready()},
            "sentiment": {"name": sentiment.model_name(), "ready": sentiment.is_ready()},
            "anonymizer": {"backend": anonymizer.backend_name()},
        },
    }


@app.get("/models")
def list_models():
    """Default, suggested and currently warm-cached models, per pipeline step."""
    cached = model_registry.cached_models()
    return {
        "sentiment": {
            "default": sentiment.DEFAULT_MODEL_NAME,
            "suggested": SUGGESTED_MODELS["sentiment"],
            "cached": cached.get("sentiment-analysis", []),
        },
        "ner": {
            "default": ner.DEFAULT_MODEL_NAME,
            "suggested": SUGGESTED_MODELS["ner"],
            "cached": cached.get("token-classification", []),
        },
        "topics": {
            "default": topics.DEFAULT_MODEL_NAME,
            "suggested": SUGGESTED_MODELS["topics"],
            "cached": cached.get("zero-shot-classification", []),
        },
    }


# ---------------------------------------------------------------------------
# Core analysis
# ---------------------------------------------------------------------------
@app.post("/analyze")
def analyze(payload: TextRequest):
    """Run the full pipeline on a single message and store the result."""
    if not payload.text or not payload.text.strip():
        raise HTTPException(status_code=400, detail="Text cannot be empty.")

    try:
        result = process_message(
            payload.text,
            topic_labels=payload.topic_labels,
            sentiment_model=payload.sentiment_model,
            ner_model=payload.ner_model,
            topic_model=payload.topic_model,
        )
    except RuntimeError as e:
        raise HTTPException(status_code=400, detail=str(e))

    try:
        db.save_record(result, source="single")
    except Exception as e:
        print(f"[main] Failed to store record: {e}")

    # do not return the raw text in a way that encourages storing it client-side;
    # we return both for the immediate UI, but only anonymized is persisted.
    return result


@app.post("/predict")
def predict(payload: TextRequest):
    """
    Backward-compatible sentiment-only endpoint.
    Kept so existing clients of the original API keep working.
    """
    if not payload.text or not payload.text.strip():
        raise HTTPException(status_code=400, detail="Text cannot be empty.")

    try:
        s = sentiment.analyze_sentiment(payload.text, model_id=payload.sentiment_model)
    except RuntimeError as e:
        raise HTTPException(status_code=400, detail=str(e))
    return {"text": payload.text, "label": s["label"], "score": s["score"]}


# ---------------------------------------------------------------------------
# Ingest (batch)
# ---------------------------------------------------------------------------
@app.post("/ingest", status_code=202)
async def ingest_file(
    file: UploadFile = File(...),
    sentiment_model: str | None = Form(None),
    ner_model: str | None = Form(None),
    topic_model: str | None = Form(None),
):
    """
    Start a background job that ingests a CSV or JSON file of conversations
    and runs the full pipeline on each message. Returns immediately with a
    job id; poll ``GET /ingest/status/{job_id}`` for progress and, once
    finished, the same summary this endpoint used to return directly
    (``received``, ``processed``, ``stored``, ``truncated``, ``items``).

    Parsing/validating the file happens synchronously here (fast, no model
    calls); only the actual NLP processing runs in the background.
    """
    raw = await file.read()
    if not raw:
        raise HTTPException(status_code=400, detail="Uploaded file is empty.")

    try:
        messages = ingest.parse_upload(file.filename or "", raw)
    except Exception as e:
        raise HTTPException(status_code=400, detail=f"Could not parse file: {e}")

    if not messages:
        raise HTTPException(
            status_code=400,
            detail="No messages found. Expected a 'text'/'message' column or field.",
        )

    # Cap batch size to keep the demo responsive on limited hardware.
    MAX_BATCH = 200
    truncated = len(messages) > MAX_BATCH
    messages = messages[:MAX_BATCH]

    job_id = jobs.create_job(total=len(messages))

    def run_job() -> None:
        try:
            results = process_batch(
                messages,
                sentiment_model=sentiment_model,
                ner_model=ner_model,
                topic_model=topic_model,
                on_progress=jobs.progress_callback(job_id),
            )
        except RuntimeError as e:
            jobs.fail_job(job_id, str(e))
            return
        except Exception as e:  # pragma: no cover - safety net so a job never hangs
            jobs.fail_job(job_id, f"Unexpected error: {e}")
            return

        stored = 0
        for r in results:
            try:
                db.save_record(r, source="ingest")
                stored += 1
            except Exception as e:
                print(f"[main] Failed to store ingest record: {e}")

        jobs.finish_job(job_id, {
            "received": len(messages),
            "processed": len(results),
            "stored": stored,
            "truncated": truncated,
            "items": results,
        })

    # Runs the (blocking, CPU-bound) pipeline in a plain thread rather than
    # an asyncio task, so it doesn't block the event loop and status polls
    # keep being served while it runs.
    threading.Thread(target=run_job, daemon=True).start()

    return {"job_id": job_id, "total": len(messages)}


@app.get("/ingest/status/{job_id}")
def ingest_status(job_id: str):
    """
    Progress of a background ingest job started via ``POST /ingest``.

    ``status`` is one of ``running``, ``done`` or ``error``. Once ``done``,
    ``result`` holds the same summary the old synchronous ``/ingest``
    returned directly.
    """
    job = jobs.get_job(job_id)
    if job is None:
        raise HTTPException(
            status_code=404,
            detail="Unknown job id (it may have finished long ago, or the server restarted).",
        )

    finished_at = job["finished_at"] or time.time()
    response: dict[str, Any] = {
        "status": job["status"],
        "processed": job["processed"],
        "total": job["total"],
        "elapsed_seconds": round(finished_at - job["started_at"], 1),
    }
    if job["status"] == "done":
        response["result"] = job["result"]
    elif job["status"] == "error":
        response["error"] = job["error"]
    return response


# ---------------------------------------------------------------------------
# Stored data
# ---------------------------------------------------------------------------
@app.get("/records")
def records(limit: int = 100):
    """Return recent stored records (anonymized only)."""
    return db.get_records(limit=limit)


@app.get("/records/export")
def export_records(format: str = "xml", limit: int | None = None):
    """
    Export stored (anonymized) records. ``format`` is one of the keys in
    ``db.EXPORT_FORMATS`` (currently ``xml``, ``json``, ``csv``).
    """
    try:
        content, media_type = db.export_records(format, limit=limit)
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))

    filename = f"records.{format.lower()}"
    return Response(
        content=content,
        media_type=media_type,
        headers={"Content-Disposition": f'attachment; filename="{filename}"'},
    )


@app.get("/stats")
def get_stats():
    """Aggregate statistics for the dashboard, plus the supported export formats."""
    return {**db.stats(), "export_formats": sorted(db.EXPORT_FORMATS)}


if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="0.0.0.0", port=8000)