src/knowledge_extraction/service.py

"""Pipeline facade: parsed document → candidate entries → review queue.

Mirrors the shape of `src/query/service.py` — a deterministic orchestrator over
stages that each do one thing, with the expensive step isolated and every
failure degrading rather than aborting.

Cost discipline, carried from the prototype and worth keeping: **dry-run, then a
small pilot, then the full run.** A dry run makes zero API calls and prints the
token estimate, so the bill is knowable before it is incurred.
"""

from __future__ import annotations

import time

from ..middlewares.logging import get_logger
from .cluster import cluster_mentions
from .cluster.normalize import normalize
from .diff import diff_glossary
from .extract import (
    build_glossary_prompt,
    est_tokens,
    extract_formula,
    extract_glossary,
    extract_rule,
    extract_summary,
    prefix_tokens,
)
from .filters import abbrev_pairs, extract_mentions, rule_candidates
from .ids import term_id
from .link import link_all
from .models import (
    CallUsage,
    Chunk,
    ClusterResult,
    FilterResult,
    ParsedDoc,
    RejectedField,
)
from .queue import build_queue
from .rank import rank_evidence, top_k
from .settings import EVIDENCE_K
from .validate import (
    evidence_text,
    find_conflicts,
    latex_source,
    rounds_available,
    validate_entry,
)

logger = get_logger("knowledge_extraction")


class ExtractionResult:
    def __init__(self) -> None:
        self.glossary: list[dict] = []
        self.rules: list[dict] = []
        self.formulas: list[dict] = []
        self.brief: dict | None = None
        self.review_queue: list[dict] = []
        self.rejected: list[RejectedField] = []
        # i2: clusters answered from an approved entry instead of a paid
        # call. Reported rather than inferred - a branch that skips work
        # silently looks identical to one that found nothing.
        self.reused: int = 0
        self.usages: list[CallUsage] = []
        # Dangle counts from the linking stage. Reported rather than swallowed:
        # a link stage that resolves nothing looks identical to one that works.
        self.links: dict[str, int] = {}

    @property
    def total_tokens(self) -> tuple[int, int, int]:
        return (
            sum(u.prompt_tokens for u in self.usages),
            sum(u.cached_tokens for u in self.usages),
            sum(u.completion_tokens for u in self.usages),
        )


def run_filters(doc: ParsedDoc, use_span_filter: bool = True) -> FilterResult:
    """All free stages. Zero API calls."""
    pairs = abbrev_pairs(doc.chunks)
    mentions = extract_mentions(doc.chunks) if use_span_filter else []
    return FilterResult(
        doc_id=doc.doc_id,
        mentions=mentions,
        rule_candidates=rule_candidates(doc.chunks),
        abbrev_pairs=pairs,
    )


def build_clusters(doc: ParsedDoc, filtered: FilterResult) -> ClusterResult:
    clustered = cluster_mentions(filtered.mentions, filtered.abbrev_pairs, doc.doc_id)
    rank_evidence(clustered.clusters, doc.chunks)
    return clustered


def estimate_cost(
    doc: ParsedDoc, clustered: ClusterResult, filtered: FilterResult, limit: int | None = None
) -> dict:
    """Dry run: exact prompts are built, nothing is sent.

    **`estimated_prompt_tokens` covers all four branches as of 2026-09-11.** It
    used to sum the glossary branch alone while carrying a name that reads as a
    total, so every consumer — including the spend ceiling that gates on it —
    was reading a partial number as a whole one. The glossary figure stays exact
    (its prompts are really built); the other three are approximated from the
    cached prefix plus the chunk text they are handed, which undercounts the
    evidence block's own framing by a few percent. Both are labelled, because a
    ceiling that does not know which half of its input is a guess is not a
    ceiling.

    `clusters` is reported alongside because **cluster count, not page count, is
    what drives spend** — one call per cluster — and it is known here for free.
    """
    clusters = clustered.clusters[:limit] if limit else clustered.clusters
    glossary_tokens = 0
    for cluster in clusters:
        system, user = build_glossary_prompt(cluster, doc.chunks)
        glossary_tokens += est_tokens(system) + est_tokens(user)

    by_id = {c.chunk_id: c for c in doc.chunks}
    rule_chunks = [
        by_id[c.chunk_id]
        for c in (filtered.rule_candidates[:limit] if limit else filtered.rule_candidates)
        if c.chunk_id in by_id
    ]
    formula_chunks = [c for c in doc.chunks if c.has_formula]
    formula_chunks = formula_chunks[:limit] if limit else formula_chunks

    rule_tokens = sum(prefix_tokens("rule") + est_tokens(c.text or "") for c in rule_chunks)
    formula_tokens = sum(
        prefix_tokens("formula") + est_tokens(c.text or "") for c in formula_chunks
    )
    # The summary branch is handed the WHOLE document — few calls, but roughly a
    # quarter of all input tokens, and it cannot be filtered (X7).
    summary_tokens = prefix_tokens("summary") + sum(
        est_tokens(c.text or "") for c in doc.chunks
    )

    by_branch = {
        "glossary": glossary_tokens,
        "rule": rule_tokens,
        "formula": formula_tokens,
        "summary": summary_tokens,
    }
    return {
        "clusters": len(clusters),
        "glossary_calls": len(clusters),
        "rule_calls": len(rule_chunks),
        "formula_calls": len(formula_chunks),
        "summary_calls": 1,
        "total_calls": len(clusters) + len(rule_chunks) + len(formula_chunks) + 1,
        "estimated_prompt_tokens": sum(by_branch.values()),
        "estimated_prompt_tokens_by_branch": by_branch,
        "glossary_tokens_are_exact": True,
        "note": (
            "estimate only — real counts come from the API usage object. Glossary "
            "prompts are built for real; rule/formula/summary are approximated from "
            "prefix + chunk text and run slightly low."
        ),
    }


def extract_all(
    doc: ParsedDoc,
    clustered: ClusterResult,
    filtered: FilterResult,
    extractor,
    limit: int | None = None,
    active_glossary: list[dict] | None = None,
    branches: tuple[str, ...] = ("glossary", "rule", "formula", "summary"),
) -> ExtractionResult:
    """The paid stage plus validation, diff and queue."""
    out = ExtractionResult()
    started = time.time()

    if "glossary" in branches:
        _run_glossary(
            doc, clustered, extractor, out, limit, _approved_index(active_glossary)
        )
    if "rule" in branches:
        _run_rules(doc, filtered, extractor, out, limit)
    if "formula" in branches:
        _run_formulas(doc, extractor, out, limit)
    if "summary" in branches:
        _run_summary(doc, extractor, out)

    # Links resolve only once every branch has produced entries — a reference
    # cannot be made to something that does not exist yet.
    out.links = link_all(out.glossary, out.rules, out.formulas, clustered, out.brief)

    out.glossary = diff_glossary(out.glossary, active_glossary or [])

    # Last, deliberately: these are aggregates over the OTHER branches, so they
    # can only be correct once every branch has finished and the diff has run.
    # Computing them inside `_run_summary` would quietly depend on the summary
    # branch executing after the other three, which is true today and is not a
    # property anything guarantees.
    _finalise_domain(out)

    out.review_queue = build_queue(out.glossary, out.brief)

    prompt, cached, completion = out.total_tokens
    logger.info(
        "extraction complete",
        doc_id=doc.doc_id,
        glossary=len(out.glossary),
        rules=len(out.rules),
        formulas=len(out.formulas),
        rejected_fields=len(out.rejected),
        calls=len(out.usages),
        **out.links,
        prompt_tokens=prompt,
        cached_tokens=cached,
        completion_tokens=completion,
        seconds=round(time.time() - started, 1),
    )
    return out


def _approved_index(active_glossary: list[dict] | None) -> dict[str, dict]:
    """Canonical key -> the approved payload, for recognising a known term (i2).

    Keyed with the SAME normaliser that buckets mentions inside a document, so
    one rule decides sameness at both scales rather than two that can drift.
    Both `term` and `full_name` are indexed because a later document may write
    the expansion where the first wrote the abbreviation.
    """
    index: dict[str, dict] = {}
    for payload in active_glossary or []:
        for surface in (payload.get("term"), payload.get("full_name")):
            key = normalize(surface or "")
            if key:
                index.setdefault(key, payload)
    return index


def _reuse_approved(approved: dict, doc: ParsedDoc, cluster) -> dict:
    """Carry an approved definition onto this document without paying for it.

    The entry is this document's — `term_id` is re-minted from *its* `doc_id`,
    because ids stay document-scoped (decision (a), 2026-09-14) and sharing one
    id across documents would move every existing `entity_id` and orphan the
    approvals this whole scheme protects.

    **Provenance is deliberately left pointing at the document the definition
    actually came from.** It would be easy to rewrite it to this document and
    wrong: nothing was extracted here, so there is no span in this document to
    cite, and a fabricated one is exactly the unfalsifiable claim the span check
    exists to refuse. Which documents re-mention the term is what `doc_ids`
    (m3) will record once it lands.
    """
    entry = dict(approved)
    entry["term_id"] = term_id(doc.doc_id, cluster.canonical)
    entry["mention_count"] = getattr(cluster, "mention_count", 0)
    entry["diff_status"] = "duplicate"
    return entry


def _run_glossary(doc, clustered, extractor, out, limit, approved=None) -> None:
    clusters = clustered.clusters[:limit] if limit else clustered.clusters
    approved = approved or {}
    for cluster in clusters:
        # i2: a term this company has already approved is not re-extracted. One
        # LLM call per cluster is the whole cost of this branch, so every term
        # recognised here is a call not made — and it is also what stops the
        # same concept arriving as `term 1.1` and `term 1.2`.
        reused = approved.get(normalize(cluster.canonical))
        if reused is not None:
            out.glossary.append(_reuse_approved(reused, doc, cluster))
            out.reused += 1
            continue

        entry = None
        max_round = rounds_available(cluster, EVIDENCE_K)

        for round_index in range(max_round + 1):
            entry, usage = extract_glossary(
                cluster, doc.chunks, extractor, doc.doc_id, EVIDENCE_K, round_index
            )
            out.usages.append(usage)
            if entry is None:
                continue

            source = evidence_text(
                # `doc.chunks` passed so a linked partner joins the window here
                # exactly as it does in the prompt (`_evidence_for`). If these two
                # ever diverge, a quote from a partner chunk is nulled as
                # unlocatable and the call is wasted.
                top_k(cluster, EVIDENCE_K, round_index, doc.chunks), doc.chunks
            )
            entry, rejections = validate_entry(entry, "glossary", source, cluster.canonical)
            out.rejected.extend(rejections)

            if entry.definition:
                if round_index > 0:
                    entry.extraction_status = "escalated"
                break
            # Null definition -> escalate to the next K chunks.

        if entry is None:
            continue
        if not entry.definition:
            entry.extraction_status = "no_definition_found"

        conflicting, variants = find_conflicts(
            [entry.definition] if entry.definition else []
        )
        entry.definition_conflict = conflicting
        entry.conflict_variants = variants
        out.glossary.append(entry.model_dump(mode="json"))


def _run_rules(doc, filtered, extractor, out, limit) -> None:
    by_id: dict[str, Chunk] = {c.chunk_id: c for c in doc.chunks}
    candidates = filtered.rule_candidates[:limit] if limit else filtered.rule_candidates
    seen: set[str] = set()
    for candidate in candidates:
        chunk = by_id.get(candidate.chunk_id)
        if chunk is None:
            continue
        entry, usage = extract_rule(candidate, chunk, extractor, doc.doc_id)
        out.usages.append(usage)
        if entry is None:
            continue
        label = (entry.statement or entry.rule_id)[:60]
        entry, rejections = validate_entry(entry, "rule", chunk.text, label)
        out.rejected.extend(rejections)
        key = (entry.statement or "").strip().casefold()
        if key and key in seen:
            continue
        if key:
            seen.add(key)
        out.rules.append(entry.model_dump(mode="json"))


def _run_formulas(doc, extractor, out, limit) -> None:
    chunks = [c for c in doc.chunks if c.has_formula]
    chunks = chunks[:limit] if limit else chunks
    seen: set[str] = set()
    for chunk in chunks:
        entry, usage = extract_formula(chunk, extractor, doc.doc_id)
        out.usages.append(usage)
        if entry is None:
            continue
        entry, rejections = validate_entry(
            entry,
            "formula",
            chunk.text,
            entry.name or chunk.chunk_id,
            latex_text=latex_source([chunk]),
        )
        out.rejected.extend(rejections)
        key = (entry.formula_latex or "").strip()
        if key and key in seen:
            continue
        if key:
            seen.add(key)
        out.formulas.append(entry.model_dump(mode="json"))


def _finalise_domain(out: ExtractionResult) -> None:
    """Fill the domain entry's aggregate fields. Deterministic, no LLM, no spend.

    `subdomains` is rolled up from the glossary entries' own `subdomain_tags`
    rather than asked of the model. Each tag was already chosen once, per term,
    with that term's evidence in front of the model; a second document-level
    classification would be the same judgement made with less context and no
    way to check it. Ordered by tag frequency, descending, with an alphabetical
    tiebreak so two runs of one document agree.

    A missing domain entry is not an error — the summary branch may not have
    run, or may have failed schema validation, and the other three artifacts
    stand on their own.
    """
    if out.brief is None:
        return

    counts: dict[str, int] = {}
    for entry in out.glossary:
        for tag in entry.get("subdomain_tags") or []:
            counts[tag] = counts.get(tag, 0) + 1

    out.brief["subdomains"] = [
        tag for tag, _ in sorted(counts.items(), key=lambda kv: (-kv[1], kv[0]))
    ]
    out.brief["n_terms"] = len(out.glossary)
    out.brief["n_formulas"] = len(out.formulas)
    out.brief["n_rules"] = len(out.rules)


def _document_outline(chunks) -> list[str]:
    """The document's own heading hierarchy, in reading order, de-duplicated.

    Derived, never generated. Chunks are walked in order and each breadcrumb is
    emitted outermost-first the first time it is seen, so a section heading
    appears once at the position it first occurs rather than once per chunk
    beneath it.

    This is the cheapest domain knowledge the pipeline has: it is verbatim
    source structure, so unlike every other field on this branch there is
    nothing here a model could get wrong. Documents with no headings yield an
    empty list, which is a fact about the document, not a failure.
    """
    outline: list[str] = []
    seen: set[str] = set()
    for chunk in chunks:
        for heading in chunk.heading_path or []:
            if heading and heading not in seen:
                seen.add(heading)
                outline.append(heading)
    return outline


def _run_summary(doc, extractor, out) -> None:
    entry, usage = extract_summary(doc.chunks, extractor, doc.doc_id)
    out.usages.append(usage)
    if entry is None:
        return
    # `outline` is set here rather than in the branch: it is read off the
    # artifact, so it needs no model and no validation.
    entry.outline = _document_outline(doc.chunks)

    # Span-checked as of 2026-09-02. The branch used to be exempt because it
    # asked the model to SUMMARISE, and a summary cannot be located in the
    # source. It now asks the model to LOCATE, so `title` and `purpose_verbatim`
    # are checkable like any other transcription — a value that is not verbatim
    # is nulled and recorded in `rejected`, never repaired.
    #
    # The haystack is the WHOLE document: unlike the other branches this one is
    # given every chunk, so a quote may legitimately come from anywhere in it.
    source = evidence_text([c.chunk_id for c in doc.chunks], doc.chunks)
    entry, rejections = validate_entry(entry, "summary", source, entry.brief_id)
    out.rejected.extend(rejections)

    out.brief = entry.model_dump(mode="json")