src/graphtrace/extraction.py

"""Per-chunk knowledge-graph extraction with verbatim provenance.

Each chunk is sent to the extractor model once. Every returned triple or result must
carry a quote that anchors verbatim in that chunk (``quotes.locate``); anything else is
rejected with a recorded reason. Accepted records keep the chunk ID and character
offsets of their evidence. Raw extractor replies live in the LLM cache; this module
writes one JSONL record per chunk with accepted and rejected items, so precision of the
extraction stage can be audited and measured.
"""

from __future__ import annotations

import json
from collections import Counter
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

from .canonical import is_generic
from .llm import LLMClient, LLMError, UsageMeter, call_structured
from .prompts import EXTRACT_SYSTEM, THIS_PAPER, ExtractionOut, extract_user
from .quotes import locate
from .schemas import SIGNATURES, Chunk, Document
from .verification import strip_headings


def _validate_triple(t, chunk_text: str) -> tuple[dict | None, str | None]:
    subj_type = "Paper" if t.subject.strip() == THIS_PAPER else t.subject_type
    obj_type = t.object_type
    if t.object.strip() == THIS_PAPER:
        return None, "paper_as_object"
    if (subj_type, obj_type) not in SIGNATURES[t.relation]:
        return None, "signature"
    if subj_type == "Paper" and t.subject.strip() != THIS_PAPER:
        return None, "foreign_paper_subject"
    if subj_type != "Paper" and is_generic(t.subject):
        return None, "generic_name"
    if is_generic(t.object):
        return None, "generic_name"
    loc = locate(t.quote, chunk_text)
    if loc is None:
        return None, "unanchored_quote"
    return {
        "subject": t.subject.strip(),
        "subject_type": subj_type,
        "relation": t.relation,
        "object": t.object.strip(),
        "object_type": obj_type,
        "quote": t.quote,
        "start_char": loc[0],
        "end_char": loc[1],
        "elided": loc[2],
    }, None


def _validate_result(r, chunk_text: str) -> tuple[dict | None, str | None]:
    # THIS_PAPER as the method is resolved to the paper's proposed method at graph build
    # (or dropped there); it is never allowed as a dataset or metric.
    if (r.method.strip() != THIS_PAPER and is_generic(r.method)) or is_generic(r.dataset) or is_generic(r.metric):
        return None, "generic_name"
    if not any(ch.isdigit() for ch in r.value):
        return None, "non_numeric_value"
    loc = locate(r.quote, chunk_text)
    if loc is None:
        return None, "unanchored_quote"
    if r.value.strip().rstrip("%") not in chunk_text[loc[0] : loc[1]]:
        return None, "value_not_in_quote"
    return {
        "method": r.method.strip(),
        "dataset": r.dataset.strip(),
        "metric": r.metric.strip(),
        "value": r.value.strip(),
        "quote": r.quote,
        "start_char": loc[0],
        "end_char": loc[1],
        "elided": loc[2],
    }, None


def extract_chunk(client: LLMClient, doc: Document, chunk: Chunk, meter: UsageMeter | None = None) -> dict:
    """One chunk -> record with accepted/rejected triples and results."""
    # Quotes are anchored against the chunk text exactly as stored (headings included).
    rec: dict = {"chunk_id": chunk.chunk_id, "doc_id": doc.doc_id, "extractor_model": client.spec.model}
    try:
        out = call_structured(
            client, meter, "extract", EXTRACT_SYSTEM, extract_user(doc.title, chunk.section_heading, strip_headings(chunk.text)), ExtractionOut, max_tokens=4096
        )
    except LLMError as e:
        rec.update(error=str(e), triples=[], results=[], rejected=[])
        return rec
    triples, results, rejected = [], [], []
    for t in out.triples:
        ok, why = _validate_triple(t, chunk.text)
        (triples.append(ok) if ok else rejected.append({"kind": "triple", "reason": why, **t.model_dump()}))
    for r in out.results:
        ok, why = _validate_result(r, chunk.text)
        (results.append(ok) if ok else rejected.append({"kind": "result", "reason": why, **r.model_dump()}))
    rec.update(error=None, triples=triples, results=results, rejected=rejected)
    return rec


def extract_corpus(
    client: LLMClient,
    docs: dict[str, Document],
    chunks: list[Chunk],
    out_path: str | Path,
    workers: int = 4,
) -> dict:
    """Resumable: chunks already present in ``out_path`` are skipped."""
    out_path = Path(out_path)
    out_path.parent.mkdir(parents=True, exist_ok=True)
    done: set[str] = set()
    if out_path.exists():
        for line in out_path.read_text(encoding="utf-8").splitlines():
            if line.strip():
                r = json.loads(line)
                if not r.get("error"):
                    done.add(r["chunk_id"])
    todo = [c for c in chunks if c.chunk_id not in done]
    meter = UsageMeter()
    with out_path.open("a", encoding="utf-8") as f, ThreadPoolExecutor(max_workers=workers) as pool:
        for rec in pool.map(lambda c: extract_chunk(client, docs[c.doc_id], c, meter), todo):
            f.write(json.dumps(rec, ensure_ascii=False) + "\n")
            f.flush()
    return {"chunks_total": len(chunks), "chunks_skipped": len(done), "chunks_processed": len(todo), "usage": meter.summary()}


def extraction_stats(records: list[dict]) -> dict:
    reasons = Counter(r["reason"] for rec in records for r in rec.get("rejected", []))
    accepted_t = sum(len(rec.get("triples", [])) for rec in records)
    accepted_r = sum(len(rec.get("results", [])) for rec in records)
    rejected = sum(reasons.values())
    proposed = accepted_t + accepted_r + rejected
    return {
        "chunks": len(records),
        "chunk_errors": sum(1 for r in records if r.get("error")),
        "triples_accepted": accepted_t,
        "results_accepted": accepted_r,
        "items_rejected": rejected,
        "rejection_reasons": dict(reasons),
        "unanchored_rate": (reasons.get("unanchored_quote", 0) / proposed) if proposed else None,
        "relations": dict(Counter(t["relation"] for rec in records for t in rec.get("triples", []))),
    }