src/graphtrace/cli.py

"""Command line: ``uv run graphtrace --help``."""

from __future__ import annotations

import json
from pathlib import Path

import typer
from rich.console import Console
from rich.table import Table

from .config import ExperimentConfig, load_project, load_yaml, make_client, write_json

app = typer.Typer(add_completion=False, no_args_is_help=True, help="GraphTrace-Bench: provenance-linked KG QA benchmark.")
console = Console()
CONFIG = typer.Option("configs/default.yaml", "--config", "-c", help="Project config YAML.")


def _exp(path: str) -> ExperimentConfig:
    return ExperimentConfig(**load_yaml(path))


@app.command()
def corpus(config: str = CONFIG) -> None:
    """Select the corpus from QASPER, chunk it, write data/processed/."""
    from .data.qasper import download
    from .pipeline import prepare_corpus

    cfg = load_project(config)
    download(cfg.paths.raw, cfg.corpus.splits)
    stats = prepare_corpus(cfg)
    console.print_json(json.dumps({k: v for k, v in stats.items() if k not in ("selection_order", "paper_datasets")}, default=str))


@app.command("demo-questions")
def demo_questions(config: str = CONFIG) -> None:
    """Select the demo's predefined QASPER questions (rule in graphtrace.demo) into data/processed/."""
    from .data.qasper import load_split
    from .demo import select_demo_questions
    from .pipeline import load_corpus, write_jsonl

    cfg = load_project(config)
    papers, chunks = load_corpus(cfg)
    questions = []
    for split in cfg.corpus.splits:
        questions.extend(load_split(cfg.paths.raw, split)[1])
    rows = select_demo_questions(questions, chunks, {d: p.title for d, p in papers.items()})
    write_jsonl(Path(cfg.paths.processed) / "demo_questions.jsonl", rows)
    console.print(f"{len(rows)} demo questions written to {Path(cfg.paths.processed) / 'demo_questions.jsonl'}")


@app.command("estimate-extraction")
def estimate_extraction(config: str = CONFIG, output_tokens_per_chunk: int = 400) -> None:
    """Rough cost of KG extraction (one call per chunk) before spending anything."""
    from .config import load_models
    from .pipeline import load_corpus
    from .prompts import EXTRACT_SYSTEM
    from .tokens import count_tokens

    cfg = load_project(config)
    _, chunks = load_corpus(cfg)
    spec = load_models(cfg.models_file)[cfg.extraction.model]
    tin = sum(c.n_tokens + count_tokens(EXTRACT_SYSTEM) + 60 for c in chunks)
    tout = len(chunks) * output_tokens_per_chunk
    usd = (tin * spec.price_in_per_mtok + tout * spec.price_out_per_mtok) / 1e6
    console.print_json(json.dumps({"model": spec.model, "calls": len(chunks), "input_tokens": tin, "output_tokens_assumed": tout, "usd": round(usd, 2)}))


@app.command()
def extract(config: str = CONFIG, limit: int = typer.Option(0, help="Only the first N chunks (0 = all).")) -> None:
    """Run LLM extraction over every chunk (resumable; raw calls cached in cache/llm)."""
    from .extraction import extract_corpus, extraction_stats
    from .pipeline import load_corpus, read_jsonl

    cfg = load_project(config)
    papers, chunks = load_corpus(cfg)
    if limit:
        chunks = chunks[:limit]
    client = make_client(cfg.extraction.model, cfg)
    out = Path(cfg.paths.graph) / "extractions.jsonl"
    summary = extract_corpus(client, papers, chunks, out, workers=cfg.extraction.workers)
    stats = extraction_stats(read_jsonl(out))
    write_json(Path(cfg.paths.graph) / "extraction_stats.json", {**summary, **stats})
    console.print_json(json.dumps({**summary, **stats}, default=str))


@app.command("build-graph")
def build_graph(config: str = CONFIG, nli: bool = typer.Option(True, help="Score edge support with the NLI model.")) -> None:
    """Canonicalise entities and write nodes/edges/results JSONL."""
    from .pipeline import build

    console.print_json(json.dumps(build(load_project(config), score_confidence=nli), default=str))


@app.command("graph-stats")
def graph_stats(config: str = CONFIG) -> None:
    from .graph import KG

    console.print_json(json.dumps(KG.load(load_project(config).paths.graph).stats(), default=str))


@app.command("generate-bench")
def generate_bench(config: str = CONFIG, paraphrase: bool = typer.Option(False, help="Paraphrase templated questions with the LLM.")) -> None:
    """Generate candidate items (all start unreviewed) into data/bench/candidates.jsonl."""
    from collections import defaultdict

    from .bench import Generator, write_items
    from .bench import paraphrase as para
    from .graph import KG
    from .pipeline import load_corpus

    cfg = load_project(config)
    _, chunks = load_corpus(cfg)
    by_doc = defaultdict(list)
    for c in chunks:
        by_doc[c.doc_id].append(c)
    kg = KG.load(cfg.paths.graph)
    items, stats = Generator(kg, dict(by_doc), cfg.bench).generate()
    if paraphrase:
        if not cfg.paraphrase_model:
            raise typer.BadParameter("set paraphrase_model in the project config")
        client = make_client(cfg.paraphrase_model, cfg)
        rejected = 0
        new = []
        for it in items:
            names = [kg.name(s.subject) for s in it.required_path[:1]] if it.category != "single_hop" else []
            q, why = para(client, it, [n for n in names if n in it.question_template])
            rejected += why is not None
            new.append(it.model_copy(update={"question": q}))
        items = new
        stats["paraphrase_rejected"] = rejected
    write_items(Path(cfg.paths.bench) / "candidates.jsonl", items)
    write_json(Path(cfg.paths.bench) / "generation_stats.json", {**stats, "bench_config": cfg.bench.model_dump(), "total": len(items)})
    console.print_json(json.dumps({**stats, "total": len(items)}))


@app.command()
def review(config: str = CONFIG, reviewer: str = typer.Option(..., help="Your name or initials, stored with each decision."), category: str = "") -> None:
    """Interactively review unreviewed items against their source evidence."""
    from .bench import append_review, apply_reviews, make_review, read_items, read_reviews
    from .graph import KG
    from .pipeline import load_corpus

    cfg = load_project(config)
    bench = Path(cfg.paths.bench)
    from .linking import EntityLinker

    kg = KG.load(cfg.paths.graph, cfg.bench.min_edge_confidence)
    _, chunks = load_corpus(cfg)
    ctext = {c.chunk_id: c.text for c in chunks}
    items = apply_reviews(read_items(bench / "candidates.jsonl"), read_reviews(bench / "reviews.jsonl"), EntityLinker(kg))
    todo = [i for i in items if i.review_status == "unreviewed" and (not category or i.category == category)]
    console.print(f"{len(todo)} items to review. Keys: [a]ccept [r]eject [e]dit [s]kip [q]uit")
    for it in todo:
        console.rule(f"{it.item_id}  ({it.category}, {it.template_id}, {it.difficulty} hop)")
        console.print(f"[bold]Q:[/bold] {it.question}")
        console.print(f"[bold]Gold decision:[/bold] {it.gold_decision}   [bold]Gold answer:[/bold] {it.gold_answer}")
        console.print(f"[dim]{it.annotation_notes}[/dim]")
        for s in it.required_path:
            if s.edge_id in kg.edges:
                e = kg.edges[s.edge_id]
                console.print(f"  ({kg.name(e.subject)}) -[{e.relation}]-> ({kg.name(e.object)})  [{e.source_chunk_id}]")
                console.print(f'    "{ctext.get(e.source_chunk_id, "")[e.start_char : e.end_char]}"')
            elif s.edge_id in kg.results:
                r = kg.results[s.edge_id]
                console.print(f"  result {kg.name(r.method)} / {kg.name(r.dataset)} / {kg.name(r.metric)} = {r.value} [{r.source_chunk_id}]")
                console.print(f'    "{ctext.get(r.source_chunk_id, "")[r.start_char : r.end_char]}"')
        choice = typer.prompt("decision", default="s").strip().lower()
        if choice == "q":
            break
        if choice == "s":
            continue
        notes = typer.prompt("notes", default="")
        if choice == "a":
            append_review(bench / "reviews.jsonl", make_review(it.item_id, "accepted", reviewer, notes))
        elif choice == "r":
            append_review(bench / "reviews.jsonl", make_review(it.item_id, "rejected", reviewer, notes))
        elif choice == "e":
            q = typer.prompt("question", default=it.question)
            dec = typer.prompt("gold decision (answer/abstain/flag_conflict)", default=it.gold_decision)
            ans = typer.prompt("gold answer (separate with ;)", default="; ".join(it.gold_answer))
            append_review(
                bench / "reviews.jsonl",
                make_review(it.item_id, "edited", reviewer, notes, question=q, gold_decision=dec, gold_answer=[a.strip() for a in ans.split(";") if a.strip()]),
            )


@app.command("freeze")
def freeze_cmd(config: str = CONFIG, version: str = "v1") -> None:
    """Export only human-accepted/edited items as the release benchmark."""
    from .bench import freeze, read_items, read_reviews, write_items
    from .graph import KG
    from .linking import EntityLinker

    cfg = load_project(config)
    bench = Path(cfg.paths.bench)
    linker = EntityLinker(KG.load(cfg.paths.graph, cfg.bench.min_edge_confidence))
    release, man = freeze(read_items(bench / "candidates.jsonl"), read_reviews(bench / "reviews.jsonl"), linker)
    write_items(bench / "release" / f"benchmark_{version}.jsonl", release)
    write_json(bench / "release" / f"benchmark_{version}.manifest.json", man)
    console.print_json(json.dumps(man))


@app.command()
def run(experiment: str, config: str = CONFIG, offline: bool = typer.Option(False, help="Replay cached calls only; fail on a miss.")) -> None:
    """Run all methods of an experiment over the benchmark (resumable)."""
    from .runner import run_experiment

    d = run_experiment(load_project(config), _exp(experiment), offline=offline)
    console.print(f"predictions written to {d / 'predictions.jsonl'}")


@app.command()
def evaluate(
    experiment: str,
    config: str = CONFIG,
    judge: bool = typer.Option(True, "--judge/--no-judge", help="Call the evaluation judge for non-gold cited chunks (default on)."),
) -> None:
    from .runner import evaluate_run

    m = evaluate_run(load_project(config), _exp(experiment), run_judge=judge)
    t = Table("method", "n (headline)", "errors", "accuracy", "path valid", "citation rate", "grounding", "false-premise rej.")
    for name, agg in m["headline"].items():
        t.add_row(
            name,
            str(agg["n"]),
            str(agg["errors"]),
            *(f"{agg[k]['mean']:.3f}" for k in ("correct", "path_all_valid", "citation_rate", "grounding_precision", "false_premise_rejected")),
        )
    console.print(t)
    console.print(f"evaluation judge: {m['judge_status']}")


@app.command()
def report(experiment: str, config: str = CONFIG, plots: bool = True) -> None:
    from .report import write_report

    cfg = load_project(config)
    console.print(f"wrote {write_report(Path(cfg.paths.runs) / _exp(experiment).name, with_plots=plots)}")


@app.command("estimate-cost")
def estimate_cost_cmd(experiment: str, config: str = CONFIG, n_items: int = 0) -> None:
    """Estimate the answer-stage cost of an experiment before running it."""
    from .bench import read_items
    from .runner import estimate_cost

    cfg = load_project(config)
    exp = _exp(experiment)
    n = n_items or len(read_items(exp.bench_file))
    console.print_json(json.dumps(estimate_cost(exp, cfg, n)))


@app.command()
def serve(config: str = CONFIG, host: str = "127.0.0.1", port: int = 8000) -> None:
    """Serve the JSON API (and the demo UI if built)."""
    import uvicorn

    from .app.server import create_app

    uvicorn.run(create_app(load_project(config)), host=host, port=port)


if __name__ == "__main__":
    app()