src/knowledge_extraction/cli.py

"""Offline runner for the knowledge-extraction pipeline.

The pipeline runs a few times a year, triggered by an admin — so a script over a
parsed artifact is the honest entry point, and any HTTP surface is a convenience
layer over this, never the other way round.

Takes a **parsed-document artifact**, never a PDF: extraction does not parse.
Every stage writes its own JSON so a later stage can be re-run without repeating
an earlier one, which matters because prompt iteration is the main development
loop and the span filter is the slow part.

    # free stages only (default) — no API calls, no spend
    uv run --no-sync python -m src.knowledge_extraction.cli <artifact.json>

    # cost estimate before spending anything
    uv run --no-sync python -m src.knowledge_extraction.cli <artifact.json> --dry-run

    # small pilot, then the full run
    uv run --no-sync python -m src.knowledge_extraction.cli <artifact.json> --extract --limit 5
    uv run --no-sync python -m src.knowledge_extraction.cli <artifact.json> --extract

    # exercise the wiring with no credentials and no spend
    uv run --no-sync python -m src.knowledge_extraction.cli <artifact.json> --extract --mock

Lives inside the package rather than in `scripts/`, which is gitignored: this
runner is the pipeline's operator entry point and has to ship with the module.

**Always --dry-run before a corpus-scale run.** It builds the exact prompts,
prints the token estimate, and makes zero API calls.
"""

from __future__ import annotations

import argparse
import json
import re
import sys
from pathlib import Path

from .adapter import parsed_doc_from_artifact
from .cluster import cluster_mentions
from .extract import MockExtractor, cacheable, prefix_tokens
from .models import Mention
from .rank import rank_evidence
from .service import build_clusters, estimate_cost, extract_all, run_filters
from .settings import EVIDENCE_K

BRANCHES = ("glossary", "rule", "formula", "summary")


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description=__doc__.split("\n")[0])
    parser.add_argument("artifact", type=Path, help="parsed-document artifact JSON")
    parser.add_argument("--out-dir", type=Path, default=Path("out/knowledge"))
    parser.add_argument("--doc-id", help="override the artifact's doc_id")
    parser.add_argument("--mentions", type=Path, help="span-NER mentions JSON")
    parser.add_argument(
        "--no-span-filter",
        action="store_true",
        help="skip the span model; legend terms only (wiring check, NOT a recall run)",
    )
    parser.add_argument(
        "--extract", action="store_true", help="run the PAID extraction stage"
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="build the prompts and print a token estimate; makes no API calls",
    )
    parser.add_argument("--mock", action="store_true", help="mock extractor: no network, no spend")
    parser.add_argument("--limit", type=int, help="cap the number of items per branch (pilot)")
    parser.add_argument(
        "--branches", nargs="+", choices=BRANCHES, default=list(BRANCHES)
    )
    parser.add_argument(
        "--active-glossary", type=Path, help="approved glossary to diff against"
    )
    parser.add_argument(
        "--allow-unlabelled-artifact",
        action="store_true",
        help="run on an artifact that does not declare parser_backend (results are "
             "NOT comparable to any committed baseline)",
    )
    args = parser.parse_args(argv)

    if not args.artifact.exists():
        print(f"artifact not found: {args.artifact}", file=sys.stderr)
        return 2

    raw = json.loads(args.artifact.read_text(encoding="utf-8"))
    doc = parsed_doc_from_artifact(raw, doc_id=args.doc_id, source_ref=str(args.artifact))

    # Refuse an artifact that will not say which backend produced it. MinerU's
    # `pipeline` and `vlm` backends emit different text from the same PDF, so an
    # unlabelled artifact silently makes every number incomparable to every other
    # number — which is exactly how the two halves spent a week measuring two
    # different parses and reading the disagreement as a bug. `ParsedDocument`
    # has carried `parser_backend` all along; a bare chunk list drops it.
    if not doc.parser_backend and not args.allow_unlabelled_artifact:
        print(
            f"artifact does not declare parser_backend: {args.artifact}\n"
            "  Extraction output is only comparable to results measured on the same\n"
            "  MinerU backend. Re-export the artifact as a ParsedDocument from\n"
            "  `src.knowledge_parsing.run` (it records the backend MinerU actually\n"
            "  used, read from _middle.json), or pass --allow-unlabelled-artifact\n"
            "  to run anyway and treat the results as uncomparable.",
            file=sys.stderr,
        )
        return 2

    backend = doc.parser_backend or "UNLABELLED"
    print(
        f"[parse ] {doc.doc_id}: {len(doc.chunks)} chunks, {doc.n_pages} pages, "
        f"hash {doc.content_hash}, backend {backend}, "
        f"heading-split={doc.used_heading_split}"
    )

    filtered = run_filters(doc, use_span_filter=not (args.no_span_filter or args.mentions))
    if args.mentions:
        filtered.mentions = _load_mentions(args.mentions)
        source = "file"
    elif args.no_span_filter:
        filtered.mentions = _from_pairs(doc, filtered.abbrev_pairs)
        source = "legend stand-in (NOT a recall run)"
    else:
        source = "span filter"

    capped = sum(m.hit_span_cap for m in filtered.mentions)
    cap_note = f", {capped} hit the span cap" if capped else ""
    print(
        f"[filter] {len(filtered.abbrev_pairs)} abbreviation pairs, "
        f"{len(filtered.rule_candidates)} rule candidates"
    )
    print(f"[filter] {len(filtered.mentions)} mentions from {source}{cap_note}")

    if args.mentions or args.no_span_filter:
        clustered = cluster_mentions(filtered.mentions, filtered.abbrev_pairs, doc.doc_id)
        rank_evidence(clustered.clusters, doc.chunks)
    else:
        clustered = build_clusters(doc, filtered)
    print(
        f"[cluster] {clustered.n_mentions} mentions -> {clustered.n_clusters} clusters "
        f"(compression {clustered.compression_ratio}x)"
    )
    for cluster in clustered.clusters[:8]:
        print(
            f"         {cluster.canonical:<26} mentions={cluster.mention_count:<4} "
            f"evidence={len(cluster.evidence_chunk_ids)} "
            f"top={cluster.evidence_chunk_ids[:EVIDENCE_K]}"
        )

    args.out_dir.mkdir(parents=True, exist_ok=True)
    _dump(args.out_dir / f"{doc.doc_id}.chunks.json", doc.model_dump(mode="json"))
    _dump(args.out_dir / f"{doc.doc_id}.filters.json", filtered.model_dump(mode="json"))
    _dump(args.out_dir / f"{doc.doc_id}.clusters.json", clustered.model_dump(mode="json"))

    if args.dry_run:
        est = estimate_cost(doc, clustered, filtered, args.limit)
        print("[dry-run] NO API CALLS MADE")
        for key, value in est.items():
            print(f"          {key}: {value}")
        for branch in args.branches:
            print(
                f"          prefix[{branch}]: {prefix_tokens(branch)} tokens, "
                f"cacheable={cacheable(branch)}"
            )
        return 0

    if not args.extract:
        print(f"[write ] {args.out_dir}   (free stages only; --extract to run the paid stage)")
        return 0

    extractor = MockExtractor() if args.mock else _azure_extractor()
    if extractor is None:
        return 3
    active = (
        json.loads(args.active_glossary.read_text(encoding="utf-8"))
        if args.active_glossary
        else []
    )

    result = extract_all(
        doc,
        clustered,
        filtered,
        extractor,
        limit=args.limit,
        active_glossary=active,
        branches=tuple(args.branches),
    )

    prompt, cached, completion = result.total_tokens
    simulated = " [SIMULATED — not a quality measurement]" if args.mock else ""
    print(f"[extract] {len(result.usages)} calls{simulated}")
    print(
        f"          glossary={len(result.glossary)} rules={len(result.rules)} "
        f"formulas={len(result.formulas)}"
    )
    print(f"          tokens prompt={prompt} cached={cached} completion={completion}")
    print(f"          fields rejected by span check: {len(result.rejected)}")
    if result.links:
        # Printed, not just logged: a link stage that resolves nothing looks
        # exactly like one that works unless the dangle counts are visible.
        print(
            "[link  ] "
            + ", ".join(f"{k}={v}" for k, v in result.links.items())
        )
    no_def = sum(1 for e in result.glossary if e.get("extraction_status") == "no_definition_found")
    print(f"          abstained (no definition in document): {no_def}/{len(result.glossary)}")

    print("[queue ] top of the review queue:")
    for row in result.review_queue[:10]:
        term = str(row.get("term"))[:26]
        print(
            f"         {row['rank']:>3}. {term:<26} "
            f"n={row['mention_count']:<4} {row['review_reason']}"
        )

    _dump(args.out_dir / "glossary.json", result.glossary)
    # Named for the entity it holds. "Interpretation pack" survives in the docs
    # as the NAME OF THE DELIVERABLE - what these rules are for - but not as a
    # third identifier: rules.json / RuleEntry / kind="rule" now all agree.
    _dump(args.out_dir / "rules.json", result.rules)
    _dump(args.out_dir / "formulas.json", result.formulas)
    _dump(args.out_dir / "review_queue.json", result.review_queue)
    _dump(args.out_dir / "rejected.json", [r.model_dump(mode="json") for r in result.rejected])
    if result.brief:
        # Named for what it holds: a per-document brief. The scope-level
        # DomainContext is composed elsewhere (src/knowledge_domain/).
        _dump(args.out_dir / "document_brief.json", result.brief)
    _dump(args.out_dir / "usage.json", [u.model_dump(mode="json") for u in result.usages])
    print(f"[write ] {args.out_dir}")
    return 0


def _azure_extractor():
    from .extract import AzureExtractor

    try:
        return AzureExtractor()
    except Exception as exc:
        print(f"cannot build the Azure client: {exc}", file=sys.stderr)
        print("use --mock to exercise the pipeline without credentials", file=sys.stderr)
        return None


def _load_mentions(path: Path) -> list[Mention]:
    raw = json.loads(path.read_text(encoding="utf-8"))
    items = raw.get("mentions", raw) if isinstance(raw, dict) else raw
    return [Mention.model_validate(m) for m in items]


def _from_pairs(doc, pairs) -> list[Mention]:
    """Stand-in mentions from legend abbreviations, so the wiring is runnable
    without the span model.

    NOT a recall measurement — it only sees terms a legend block already named.

    Word-boundary matching, never substring: "PA" occurs inside "parameter",
    "pada", "capacity" and "composite", and substring matching produced 126
    spurious PA mentions on a 9-page document (77x compression instead of the
    measured 2.56x). Same trap the evidence ranker documents for headings.
    """
    surfaces = {p.abbrev for p in pairs} | {p.expansion for p in pairs}
    patterns = [
        (s, re.compile(rf"(?<!\w){re.escape(s)}(?!\w)", re.IGNORECASE)) for s in surfaces
    ]
    out: list[Mention] = []
    for chunk in doc.chunks:
        for surface, pattern in patterns:
            for match in pattern.finditer(chunk.text):
                out.append(
                    Mention(
                        surface=surface,
                        chunk_id=chunk.chunk_id,
                        char_start=match.start(),
                        char_end=match.end(),
                        label="legend",
                        score=1.0,
                    )
                )
    return out


def _dump(path: Path, payload) -> None:
    path.write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8")


if __name__ == "__main__":
    raise SystemExit(main())