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 .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, 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] = []
        self.usages: list[CallUsage] = []

    @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)

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

    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),
        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
        entry, rejections = validate_entry(entry, "rule", chunk.text, entry.rule_id)
        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
        )
        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 _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
    # Summary prose is NOT span-checked — it cannot be. Only its provenance is
    # carried, and the branch belongs on a larger tier for exactly that reason.
    out.brief = entry.model_dump(mode="json")