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 .diff import diff_glossary
from .extract import (
    build_glossary_prompt,
    est_tokens,
    extract_formula,
    extract_glossary,
    extract_rule,
    extract_summary,
)
from .filters import abbrev_pairs, extract_mentions, rule_candidates
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.domain: dict | None = None
        self.review_queue: list[dict] = []
        self.rejected: list[RejectedField] = []
        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."""
    clusters = clustered.clusters[:limit] if limit else clustered.clusters
    prompt_tokens = 0
    for cluster in clusters:
        system, user = build_glossary_prompt(cluster, doc.chunks)
        prompt_tokens += est_tokens(system) + est_tokens(user)
    return {
        "glossary_calls": len(clusters),
        "rule_calls": len(filtered.rule_candidates),
        "formula_calls": sum(1 for c in doc.chunks if c.has_formula),
        "summary_calls": 1,
        "estimated_prompt_tokens": prompt_tokens,
        "note": "estimate only — real counts come from the API usage object",
    }


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)
    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.domain)

    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.domain)

    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 _run_glossary(doc, clustered, extractor, out, limit) -> None:
    clusters = clustered.clusters[:limit] if limit else clustered.clusters
    for cluster in clusters:
        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(
                top_k(cluster, EVIDENCE_K, round_index), 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.domain 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.domain["subdomains"] = [
        tag for tag, _ in sorted(counts.items(), key=lambda kv: (-kv[1], kv[0]))
    ]
    out.domain["n_terms"] = len(out.glossary)
    out.domain["n_formulas"] = len(out.formulas)
    out.domain["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.domain_id)
    out.rejected.extend(rejections)

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