server.py

"""
Flask server for the Clinical Simplifier Web UI.

Serves the single-page front-end and exposes a streaming `/api/simplify`
endpoint that runs the full V7 pipeline (in-process, via engine.py) for one
clinical note and pushes live stage-progress events back to the browser as
Server-Sent Events, finishing with the full structured result.

Run:
    cd webapp
    python server.py
    # open http://127.0.0.1:5000
"""

from __future__ import annotations

import json
import os
import queue
import threading
import traceback
from pathlib import Path

from flask import Flask, Response, jsonify, request, send_from_directory

import engine

HERE = Path(__file__).resolve().parent
app = Flask(__name__, static_folder=str(HERE / "static"), static_url_path="/static")

# Only one pipeline run executes at a time (Engine holds a run-lock), but we
# guard the queue plumbing too so progress streams stay isolated per request.
_DEFAULT_ABLATIONS = int(os.environ.get("CC_NUM_ABLATIONS", "32"))


def _sse(stage: str, status: str, data=None) -> str:
    return "data: " + json.dumps({"stage": stage, "status": status, "data": data}) + "\n\n"


@app.route("/")
def index():
    return send_from_directory(str(HERE / "templates"), "index.html")


@app.route("/api/health")
def health():
    return jsonify({
        "status": "ok",
        "loaded": engine.ENGINE._loaded,
        "load_status": engine.ENGINE.load_status,
        "default_ablations": _DEFAULT_ABLATIONS,
        "simplifier_models": [
            {"key": m["key"], "label": m["label"], "provider": m["provider"],
             "tier": m["tier"], "default": m.get("default", False)}
            for m in engine.SIMPLIFIER_MODELS
        ],
        "novita_models": engine.NOVITA_MODELS_MENU,
        "judge_panel": engine.JUDGE_PANEL,
    })


@app.route("/api/dataset")
def dataset():
    """Return the grade-3 gold simplification triples for benchmark mode."""
    fp = engine.DATA_DIR / "laymaker_grade3_plain.json"
    try:
        rows = json.loads(fp.read_text(encoding="utf-8"))
        items = [{"id": r.get("ID"), "original": r.get("ORIGINAL"),
                  "reference": r.get("REFERENCE")} for r in rows
                 if r.get("ORIGINAL") and r.get("REFERENCE")]
        return jsonify({"ok": True, "items": items})
    except Exception as e:
        return jsonify({"ok": False, "error": str(e), "items": []})



@app.route("/api/warmup", methods=["POST"])
def warmup():
    """Pre-load glossaries + LLM so the first real request is fast."""
    def progress(stage, status, data):
        pass
    try:
        engine.ENGINE.ensure_loaded(progress)
        return jsonify({"ok": True, "load_status": engine.ENGINE.load_status})
    except Exception as e:
        traceback.print_exc()
        return jsonify({"ok": False, "error": str(e)}), 500


@app.route("/api/simplify", methods=["POST"])
def simplify():
    data = request.get_json(force=True, silent=True) or {}
    note = (data.get("note") or "").strip()
    reference = (data.get("reference") or "").strip()
    fw_model = (data.get("fw_model") or "").strip() or None
    novita_model = (data.get("novita_model") or "").strip() or None
    run_judge = bool(data.get("judge"))
    ablations = data.get("ablations")
    try:
        ablations = int(ablations) if ablations else None
    except (TypeError, ValueError):
        ablations = None

    if not note:
        return Response(_sse("error", "error", {"error": "Please enter a clinical note."}),
                        mimetype="text/event-stream")
    if len(note) > 6000:
        return Response(_sse("error", "error", {"error": "Note too long (max 6000 chars)."}),
                        mimetype="text/event-stream")

    q: "queue.Queue" = queue.Queue()
    holder: dict = {}

    def progress(stage, status, d):
        q.put({"stage": stage, "status": status, "data": d})

    def worker():
        try:
            holder["result"] = engine.ENGINE.run(
                note, progress=progress, ablations=ablations, fw_model=fw_model,
                novita_model=novita_model, reference=reference, run_judge=run_judge)
        except Exception as e:  # noqa: BLE001
            msg = str(e)
            if "503" in msg or "service_unavailable" in msg.lower() or "service unavailable" in msg.lower():
                holder["error"] = ("The selected provider is temporarily busy (503). "
                                   "Please try again in a moment, lower the ContextCite "
                                   "ablations in Settings, or switch to a Firework model "
                                   "(more reliable than the Together ones).")
            else:
                holder["error"] = f"{type(e).__name__}: {e}"
            traceback.print_exc()
        finally:
            q.put(None)

    threading.Thread(target=worker, daemon=True).start()

    def stream():
        yield _sse("queued", "running", {"detail": "starting pipeline"})
        while True:
            item = q.get()
            if item is None:
                break
            yield "data: " + json.dumps(item) + "\n\n"
        if "result" in holder:
            yield _sse("result", "done", holder["result"])
        else:
            yield _sse("error", "error", {"error": holder.get("error", "unknown error")})

    return Response(stream(), mimetype="text/event-stream",
                    headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no",
                             "Connection": "keep-alive"})


if __name__ == "__main__":
    port = int(os.environ.get("PORT", "5000"))
    host = os.environ.get("HOST", "127.0.0.1")
    print(f"\n  MedPlain UI  →  http://{host}:{port}\n")
    # threaded=True so the SSE stream + the worker thread run concurrently.
    app.run(host=host, port=port, threaded=True, debug=False)