Spaces:
Running on Zero
Running on Zero
| """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" | |
| def kg() -> KG: | |
| return KG.load(paths.graph, cfg.bench.min_edge_confidence) | |
| def linker() -> EntityLinker: | |
| return EntityLinker(kg()) | |
| def chunks() -> dict[str, Chunk]: | |
| return {c["chunk_id"]: Chunk(**c) for c in read_jsonl(Path(paths.processed) / "chunks.jsonl")} | |
| def titles() -> dict[str, str]: | |
| return {p["doc_id"]: p["title"] for p in read_jsonl(Path(paths.processed) / "papers.jsonl")} | |
| 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, | |
| } | |
| 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(), | |
| } | |
| def demo_questions() -> list[dict]: | |
| return read_jsonl(Path(paths.processed) / "demo_questions.jsonl") | |
| 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 | |
| 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() | |
| 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")) | |
| def graph_stats() -> dict: | |
| return kg().stats() | |
| 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]] | |
| 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]} | |
| 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)} | |
| 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) | |
| ] | |
| 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()]} | |
| 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()} | |
| 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 | |
| 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")) | |
| 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() | |