src/graphtrace/app/server.py

"""Read-mostly JSON API over the corpus, graph, benchmark and runs (optional extra ``app``).

Review endpoints write to the reviews JSONL file and are disabled unless the
environment variable ``GRAPHTRACE_ALLOW_REVIEW=1`` is set, so a public deployment
cannot be used to alter annotations. Static files in ``app/static`` (the demo UI) are
served at ``/`` when present.
"""

from __future__ import annotations

import importlib.util
import json
import os
import threading
from functools import lru_cache
from pathlib import Path
from typing import Literal

from fastapi import FastAPI, HTTPException, Query
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel

from ..bench import append_review, apply_reviews, make_review, read_items, read_reviews
from ..config import ProjectConfig, load_project
from ..demo import MAX_CHECK_CHUNKS, DemoInputError, LiveDemo
from ..graph import KG
from ..linking import EntityLinker
from ..paths import check_path
from ..pipeline import read_jsonl
from ..schemas import Chunk, Decision
from ..verification import NLIVerifier

STATIC = Path(__file__).parent / "static"


class ReviewIn(BaseModel):
    # Literal types make FastAPI answer 422 (not 500) for invalid values.
    status: Literal["accepted", "rejected", "edited"]
    reviewer: str
    notes: str = ""
    question: str | None = None
    gold_decision: Decision | None = None
    gold_answer: list[str] | None = None


class RetrieveIn(BaseModel):
    question: str
    gold_chunk_id: str | None = None


class CheckIn(BaseModel):
    claim: str
    chunk_ids: list[str]


def nli_available() -> bool:
    return importlib.util.find_spec("sentence_transformers") is not None


def create_app(project: ProjectConfig | None = None, nli: NLIVerifier | None = None) -> FastAPI:
    """``nli`` overrides the verifier used by the live claim check (tests pass a stub)."""
    cfg = project or load_project(os.environ.get("GRAPHTRACE_CONFIG", "configs/default.yaml"))
    paths = cfg.paths
    app = FastAPI(title="GraphTrace-Bench", version="0.1.0")
    bench_dir = Path(paths.bench)
    candidates = bench_dir / "candidates.jsonl"
    reviews_file = bench_dir / "reviews.jsonl"

    @lru_cache(maxsize=1)
    def kg() -> KG:
        return KG.load(paths.graph, cfg.bench.min_edge_confidence)

    @lru_cache(maxsize=1)
    def linker() -> EntityLinker:
        return EntityLinker(kg())

    @lru_cache(maxsize=1)
    def chunks() -> dict[str, Chunk]:
        return {c["chunk_id"]: Chunk(**c) for c in read_jsonl(Path(paths.processed) / "chunks.jsonl")}

    @lru_cache(maxsize=1)
    def titles() -> dict[str, str]:
        return {p["doc_id"]: p["title"] for p in read_jsonl(Path(paths.processed) / "papers.jsonl")}

    @lru_cache(maxsize=1)
    def demo() -> LiveDemo:
        verifier = nli if nli is not None else (NLIVerifier(cfg.extraction.nli_model) if nli_available() else None)
        return LiveDemo(chunks(), titles(), nli=verifier)

    # The NLI model is CPU-bound: one check at a time, others wait briefly and then get 503.
    nli_slot = threading.BoundedSemaphore(1)

    def items() -> dict:
        if not candidates.exists():
            return {}
        return {it.item_id: it for it in apply_reviews(read_items(candidates), read_reviews(reviews_file), linker())}

    def edge_view(eid: str) -> dict:
        g = kg()
        e = g.edges[eid]
        c = chunks().get(e.source_chunk_id)
        return {
            **e.model_dump(),
            "subject_name": g.name(e.subject),
            "object_name": g.name(e.object),
            "paper_title": titles().get(c.doc_id) if c else None,
        }

    @app.get("/api/health")
    def health() -> dict:
        return {
            "ok": True,
            "corpus": (Path(paths.processed) / "chunks.jsonl").exists(),
            "graph": (Path(paths.graph) / "edges.jsonl").exists(),
            "benchmark": candidates.exists(),
            "review_enabled": os.environ.get("GRAPHTRACE_ALLOW_REVIEW") == "1",
            "demo_questions": (Path(paths.processed) / "demo_questions.jsonl").exists(),
            "nli": nli is not None or nli_available(),
        }

    @app.get("/api/demo/questions")
    def demo_questions() -> list[dict]:
        return read_jsonl(Path(paths.processed) / "demo_questions.jsonl")

    @app.post("/api/demo/retrieve")
    def demo_retrieve(body: RetrieveIn) -> dict:
        if not chunks():
            raise HTTPException(404, "corpus not prepared")
        try:
            return demo().retrieve(body.question, body.gold_chunk_id)
        except DemoInputError as e:
            raise HTTPException(422, str(e)) from e

    @app.post("/api/demo/check")
    def demo_check(body: CheckIn) -> dict:
        d = demo()
        if len(body.chunk_ids) > MAX_CHECK_CHUNKS:
            raise HTTPException(422, f"select at most {MAX_CHECK_CHUNKS} passages")
        if d.nli is None:
            raise HTTPException(503, "the NLI verifier is not installed on this server (install the `nli` extra)")
        if not nli_slot.acquire(timeout=30):
            raise HTTPException(503, "the verifier is busy; try again in a moment")
        try:
            return d.check(body.claim, body.chunk_ids)
        except DemoInputError as e:
            raise HTTPException(422, str(e)) from e
        finally:
            nli_slot.release()

    @app.get("/api/corpus/stats")
    def corpus_stats() -> dict:
        p = Path(paths.processed) / "corpus_manifest.json"
        if not p.exists():
            raise HTTPException(404, "corpus not prepared")
        return json.loads(p.read_text(encoding="utf-8"))

    @app.get("/api/graph/stats")
    def graph_stats() -> dict:
        return kg().stats()

    @app.get("/api/graph/search")
    def search(q: str = Query(min_length=1), limit: int = 20) -> list[dict]:
        ql = q.lower()
        hits = [n for n in kg().nodes.values() if ql in n.name.lower() or any(ql in a.lower() for a in n.aliases)]
        hits.sort(key=lambda n: (-n.mentions, n.name))
        return [n.model_dump() for n in hits[:limit]]

    @app.get("/api/graph/neighbourhood")
    def neighbourhood(node_id: str, hops: int = Query(1, ge=1, le=2), limit: int = 200) -> dict:
        g = kg()
        if node_id not in g.nodes:
            raise HTTPException(404, "unknown node")
        edges = g.neighbourhood([node_id], hops)[:limit]
        node_ids = {node_id} | {e.subject for e in edges} | {e.object for e in edges}
        return {"center": node_id, "nodes": [g.nodes[n].model_dump() for n in sorted(node_ids)], "edges": [edge_view(e.edge_id) for e in edges]}

    @app.get("/api/chunks/{chunk_id}")
    def chunk(chunk_id: str) -> dict:
        c = chunks().get(chunk_id)
        if c is None:
            raise HTTPException(404, "unknown chunk")
        return {**c.model_dump(), "paper_title": titles().get(c.doc_id)}

    @app.get("/api/bench/items")
    def bench_items(status: str | None = None, category: str | None = None) -> list[dict]:
        rows = [it for it in items().values() if (status is None or it.review_status == status) and (category is None or it.category == category)]
        return [
            {
                k: v
                for k, v in it.model_dump().items()
                if k in ("item_id", "category", "template_id", "difficulty", "question", "gold_decision", "gold_answer", "review_status")
            }
            for it in sorted(rows, key=lambda i: i.item_id)
        ]

    @app.get("/api/bench/items/{item_id}")
    def bench_item(item_id: str) -> dict:
        it = items().get(item_id)
        if it is None:
            raise HTTPException(404, "unknown item")
        g = kg()
        steps = []
        for s in it.required_path:
            if s.edge_id in g.edges:
                steps.append({"kind": "edge", **edge_view(s.edge_id)})
            elif s.edge_id in g.results:
                steps.append({"kind": "result", **g.results[s.edge_id].model_dump(), "method_name": g.name(s.subject), "dataset_name": g.name(s.object)})
        return {"item": it.model_dump(), "path": steps, "chunks": [chunk(c) for c in it.gold_chunk_ids if c in chunks()]}

    @app.post("/api/bench/items/{item_id}/review")
    def review(item_id: str, body: ReviewIn) -> dict:
        if os.environ.get("GRAPHTRACE_ALLOW_REVIEW") != "1":
            raise HTTPException(403, "reviewing is disabled on this server (set GRAPHTRACE_ALLOW_REVIEW=1 locally)")
        if item_id not in items():
            raise HTTPException(404, "unknown item")
        edits = {k: v for k, v in (("question", body.question), ("gold_decision", body.gold_decision), ("gold_answer", body.gold_answer)) if v is not None}
        if edits and body.status != "edited":
            raise HTTPException(422, "edits require status 'edited'")
        r = make_review(item_id, body.status, body.reviewer, body.notes, **edits)
        append_review(reviews_file, r)
        return {"ok": True, "review": r.model_dump(), "item": items()[item_id].model_dump()}

    @app.get("/api/runs")
    def runs() -> list[dict]:
        root = Path(paths.runs)
        if not root.exists():
            return []
        return [
            {
                "name": d.name,
                "has_metrics": (d / "metrics.json").exists(),
                "smoke_test": json.loads((d / "manifest.json").read_text(encoding="utf-8")).get("smoke_test_unreviewed_items")
                if (d / "manifest.json").exists()
                else None,
            }
            for d in sorted(root.iterdir())
            if d.is_dir()
        ]

    def run_dir(name: str) -> Path:
        d = Path(paths.runs) / name
        if not d.is_dir() or "/" in name or "\\" in name or name.startswith("."):
            raise HTTPException(404, "unknown run")
        return d

    @app.get("/api/runs/{name}/metrics")
    def run_metrics(name: str) -> dict:
        p = run_dir(name) / "metrics.json"
        if not p.exists():
            raise HTTPException(404, "run not evaluated")
        return json.loads(p.read_text(encoding="utf-8"))

    @app.get("/api/runs/{name}/items/{item_id}")
    def run_item(name: str, item_id: str) -> dict:
        d = run_dir(name)
        man = d / "manifest.json"
        hashes = json.loads(man.read_text(encoding="utf-8")).get("method_config_hashes") if man.exists() else None
        # Only predictions of the run's current method configs (stale rows are ignored).
        preds = [
            r for r in read_jsonl(d / "predictions.jsonl") if r["item_id"] == item_id and (hashes is None or r.get("config_hash") == hashes.get(r["method"]))
        ]
        latest = {r["method"]: r for r in preds}
        it = items().get(item_id)
        seeds = {lk.node_id for lk in linker().link_question(it.question)} if it else set()
        out = {}
        for m, r in latest.items():
            rep = check_path(kg(), linker(), r.get("graph_path", []), seeds, set())
            out[m] = {**{k: v for k, v in r.items() if k != "trace"}, "path_check": rep.to_dict(), "verifier": r.get("trace", {}).get("verifier")}
        return {"item": it.model_dump() if it else None, "outputs": out}

    if STATIC.exists() and any(STATIC.iterdir()):
        app.mount("/", StaticFiles(directory=STATIC, html=True), name="static")
    return app


def app_factory() -> FastAPI:
    return create_app()