FabledCurator/backend/app/api/heads.py

"""Heads API (#114): train + inspect the per-concept heads that power
suggestions (replacing Camie + centroid).

POST /api/heads/train   — (re)train all eligible heads (one run at a time).
GET  /api/heads         — status: head count, last-trained, running run, the
                          per-concept head table (strength + auto-apply ready),
                          and recent training runs. The card rehydrates from
                          here so status survives navigation.
"""

from quart import Blueprint, jsonify, request
from sqlalchemy import desc, func, select

from ..extensions import get_session
from ..models import HeadAutoApplyRun, HeadTrainingRun, Tag, TagHead
from ..services.ml.heads import (
    HeadAutoApplyAlreadyRunning,
    HeadAutoApplyDisabled,
    HeadTrainingAlreadyRunning,
    start_head_auto_apply_run,
    start_head_training_run,
)

heads_bp = Blueprint("heads", __name__, url_prefix="/api/heads")


def _serialize_run(run: HeadTrainingRun) -> dict:
    return {
        "id": run.id,
        "params": run.params,
        "status": run.status,
        "started_at": run.started_at.isoformat() if run.started_at else None,
        "finished_at": run.finished_at.isoformat() if run.finished_at else None,
        "n_trained": run.n_trained,
        "n_skipped": run.n_skipped,
        "error": run.error,
    }


@heads_bp.route("/train", methods=["POST"])
async def train():
    body = await request.get_json(silent=True) or {}
    params = body.get("params") or body or {}
    async with get_session() as session:
        try:
            run_id = await session.run_sync(
                lambda s: start_head_training_run(s, params)
            )
        except HeadTrainingAlreadyRunning as running:
            return jsonify({
                "error": "training_already_running",
                "running_id": int(running.args[0]),
            }), 409
        await session.commit()
    return jsonify({"run_id": run_id, "status": "running"}), 202


@heads_bp.route("", methods=["GET"])
async def status():
    async with get_session() as session:
        count, last_trained = (
            await session.execute(
                select(func.count(), func.max(TagHead.trained_at))
            )
        ).one()
        graduated = (
            await session.execute(
                select(func.count()).where(
                    TagHead.auto_apply_threshold.is_not(None)
                )
            )
        ).scalar_one()
        running = (
            await session.execute(
                select(HeadTrainingRun.id)
                .where(HeadTrainingRun.status == "running")
                .order_by(HeadTrainingRun.id.desc())
                .limit(1)
            )
        ).scalar_one_or_none()
        runs = (
            await session.execute(
                select(HeadTrainingRun)
                .order_by(HeadTrainingRun.id.desc())
                .limit(10)
            )
        ).scalars().all()
        # The per-concept table: strongest first, capped for the admin card.
        head_rows = (
            await session.execute(
                select(
                    TagHead.tag_id, Tag.name, Tag.kind,
                    TagHead.n_pos, TagHead.n_neg, TagHead.ap,
                    TagHead.precision_cv, TagHead.recall,
                    TagHead.auto_apply_threshold, TagHead.trained_at,
                )
                .join(Tag, Tag.id == TagHead.tag_id)
                .order_by(desc(TagHead.ap))
                .limit(500)
            )
        ).all()
        heads = [
            {
                "tag_id": r.tag_id,
                "name": r.name,
                "category": r.kind.value if hasattr(r.kind, "value") else str(r.kind),
                "n_pos": r.n_pos,
                "n_neg": r.n_neg,
                "ap": r.ap,
                "precision": r.precision_cv,
                "recall": r.recall,
                "auto_apply": r.auto_apply_threshold is not None,
                "trained_at": r.trained_at.isoformat() if r.trained_at else None,
            }
            for r in head_rows
        ]
    return jsonify({
        "head_count": count,
        "graduated_count": graduated,
        "last_trained_at": last_trained.isoformat() if last_trained else None,
        "running_id": running,
        "runs": [_serialize_run(r) for r in runs],
        "heads": heads,
    })


def _serialize_apply_run(run: HeadAutoApplyRun) -> dict:
    return {
        "id": run.id,
        "dry_run": run.dry_run,
        "status": run.status,
        "started_at": run.started_at.isoformat() if run.started_at else None,
        "finished_at": run.finished_at.isoformat() if run.finished_at else None,
        "n_applied": run.n_applied,
        "report": run.report,
        "error": run.error,
    }


@heads_bp.route("/auto-apply", methods=["POST"])
async def auto_apply():
    """Trigger an earned-auto-apply sweep. {dry_run:true} previews (writes
    nothing); a real sweep needs head_auto_apply_enabled on."""
    body = await request.get_json(silent=True) or {}
    params = {"dry_run": bool(body.get("dry_run", False))}
    async with get_session() as session:
        try:
            run_id = await session.run_sync(
                lambda s: start_head_auto_apply_run(s, params)
            )
        except HeadAutoApplyAlreadyRunning as running:
            return jsonify({
                "error": "auto_apply_already_running",
                "running_id": int(running.args[0]),
            }), 409
        except HeadAutoApplyDisabled:
            return jsonify({"error": "auto_apply_disabled"}), 400
        await session.commit()
    return jsonify({"run_id": run_id, "status": "running"}), 202


@heads_bp.route("/auto-apply", methods=["GET"])
async def auto_apply_status():
    async with get_session() as session:
        running = (
            await session.execute(
                select(HeadAutoApplyRun.id)
                .where(HeadAutoApplyRun.status == "running")
                .order_by(HeadAutoApplyRun.id.desc())
                .limit(1)
            )
        ).scalar_one_or_none()
        runs = (
            await session.execute(
                select(HeadAutoApplyRun)
                .order_by(HeadAutoApplyRun.id.desc())
                .limit(10)
            )
        ).scalars().all()
    return jsonify({
        "running_id": running,
        "runs": [_serialize_apply_run(r) for r in runs],
    })