src/knowledge_parsing/contracts.py

"""Data shapes handed between pipeline stages — the parsing/extraction seam.

Extraction consumes `ParsedDocument`, never a file path and never the parser's
API. That is what keeps the parser swappable: a Tesseract or Azure Document
Intelligence path emits the same artifact, and the extraction half never learns
which one ran.

Design note — these fields were derived from real MinerU output, not designed on
paper. Verified against actual `content_list.json` output and cross-checked
against MinerU's source for both the `pipeline` and `vlm` backends.

Two things worth knowing before reviewing:

1. `bbox` and `text_level` are emitted by **both** backends — verified in
   MinerU's source, where `pipeline` and `vlm` run identical logic:

       elif para_type == BlockType.TITLE:
           title_level = get_title_level(para_block)
           if title_level != 0:
               para_content['text_level'] = title_level

   So heading structure is read from the document when MinerU detects titles,
   and only derived from numbering patterns as a fallback.

2. Heading availability is **document-dependent, not backend-dependent**. A
   document with explicit numbered headings (e.g. the BUMA standard: `2.` /
   `2.1.` / `2.1.1.`) yields a clean hierarchy. A document of mid-chapter pages
   with no headings yields none — which is why `section_no`, `heading` and
   `heading_path` stay Optional rather than required.

⚠️ `Chunk.text` must stay VERBATIM from the source document. The extraction-side
guardrail locates LLM-quoted spans literally inside this text; if the text is
ever cleaned up (lines rejoined, whitespace normalised), the lookup fails and
the field is silently set to null instead of raising. The failure looks like a
bad LLM, but the cause would be here.
"""

from __future__ import annotations

from typing import Literal

from pydantic import BaseModel, Field, computed_field

ChunkKind = Literal["text", "table", "chart", "equation", "image"]
AssetKind = Literal["figure", "chart", "diagram", "table"]

# 0.4.0 — ASSET POINTERS. A figure is no longer its own chunk: its markdown
#         reference is folded into the prose chunk at its reading-order position
#         and the figure travels as an `Asset` on that chunk. Tables stay their
#         own chunk (their linearised text earns its place as evidence) but gain
#         an `Asset` and a two-way pointer. Adds `Chunk.assets` and
#         `Chunk.referenced_by`; `Chunk.images` is kept as a DEPRECATED alias.
#
#         A MINOR bump, not a patch: chunk boundaries change, so a 0.3.x consumer
#         reading a 0.4.0 artifact would silently see a different document shape.
#         `adapter.check_schema_version` rejects a newer minor, which is the
#         designed behaviour — parsing and extraction ship together here.
# 0.3.3 — adds the computed 1-based `page_no` (S6b). Derived from `page_idx`.
# 0.3.2 — adds the optional `source_title`.
# 0.3.1 — adds the optional `normaliser_version`.
# 0.3.x were all patch bumps: purely additive, so every 0.3.0 consumer reads them
# unchanged and the new fields are null on anything written earlier.
SCHEMA_VERSION = "0.4.0"


class Asset(BaseModel):
    """A figure, chart, diagram or table referenced by a chunk. Added 0.4.0.

    **Why this exists.** Before 0.4.0 an image became its own chunk with an EMPTY
    `text`, which made it unreachable by construction: no text means no term
    mentions, no mentions means the chunk joins no cluster, and a chunk in no
    cluster is never selected as evidence. Measured across three documents: 13
    figures, every one described by a vision model we paid for, and **zero of
    them appearing in any cluster's evidence.** Meanwhile the prose that said
    "Figure 1.1" had its `![](…)` reference stripped, so nothing connected the
    sentence to the picture.

    The asset carries both halves of the provenance split, and the distinction is
    load-bearing:

    - `caption` is **verbatim document text** — printed on the page, therefore
      quotable as evidence and span-checkable like any other transcription.
    - `description` is **model-written** — a claim *about* the picture. It is
      never span-checkable and must never be quoted as evidence, which is why it
      lives here and not in `Chunk.text`.
    """

    # Content-addressed: sha256 of the asset bytes, first 16 hex chars.
    #
    # NOT the parser's own id. Mistral names images per PAGE (`img-0.jpeg`), which
    # already collides in practice — the Open Pit textbook yields 7 distinct names
    # for 8 images, so a flat store keyed on the parser's name silently loses one.
    # Content addressing also makes a re-parse idempotent and de-duplicates a
    # figure that appears in two documents.
    asset_id: str
    kind: AssetKind

    # Where the bytes live. `null` for a table (it has no bytes of its own) and
    # for a figure whose persistence failed — writing assets is a never-throw
    # seam, so a storage outage degrades to "no key" rather than failing a parse.
    storage_key: str | None = None

    page_idx: int = 0

    @computed_field  # type: ignore[prop-decorator]
    @property
    def page_no(self) -> int:
        """1-based page number, for a human and for a review UI. Derived from
        `page_idx` so the two cannot drift."""
        return self.page_idx + 1

    # `[x0, y0, x1, y1]` of the asset itself, so a UI can highlight the figure.
    # Distinct from the host chunk's own `bbox`, which stays where the TEXT is.
    bbox: list[int] | None = None

    # Verbatim, from the document's own caption block. Quotable as evidence.
    caption: str | None = None

    # ⚠️ MODEL-WRITTEN. See the class docstring. Never evidence.
    description: str | None = None
    # Which model/feature produced `description`, so a future swap is traceable
    # rather than silently changing what the field means.
    description_source: str | None = None


class Chunk(BaseModel):
    """One unit of a parsed document, ready for the extraction stage."""

    # Provenance key. Format is part of the interface: `f"{doc_id}::{ordinal:04d}"`,
    # separator `::`, agreed 2026-08-24. Every candidate entry points back here, so
    # the two halves must mint and read the same shape or provenance breaks silently.
    chunk_id: str
    doc_id: str
    kind: ChunkKind

    # Content. VERBATIM — never reflowed or normalised. See note above.
    #
    # Does NOT include the section's heading line; that lives in `heading`.
    # Consumers that need both compose them — the extraction half already builds
    # its span-check haystack as `heading + "\n" + text` and treats a heading
    # that names the term as a mention at offset 0.
    text: str

    # Location in the source document.
    #
    # Named `page_idx` deliberately: these are 0-BASED, exactly as the parser
    # reports them, with no conversion — so the artifact always matches the raw
    # parser output kept alongside it. The 1-based number a human reads is
    # `page_no` below: a review queue is about to see these, and an off-by-one
    # there is invisible until an expert opens the wrong page, so the conversion
    # is done once, here, rather than left to each consumer (S6b, 2026-09-02).
    page_idx: int
    page_idxs: list[int] = Field(default_factory=list)

    @computed_field  # type: ignore[prop-decorator]
    @property
    def page_no(self) -> int:
        """1-based page number, for humans and the review UI.

        Derived from `page_idx`, never stored, so it cannot drift from it — and
        neither normaliser has to remember to set it. `page_idx` stays 0-based
        and raw; this is the single number the reviewer sees.
        """
        return self.page_idx + 1

    # Structural context. All Optional — many documents carry no headings.
    section_no: str | None = None  # e.g. "2.1.3", when the document is numbered

    # This chunk's own section title, VERBATIM as the document writes it —
    # including wording a reader may be tempted to "fix". The BUMA standard says
    # "Physical of Availability (PA)", not "Physical Availability"; that exact
    # string must reach the extraction model, or it gets silently normalised and
    # the discrepancy never reaches the expert.
    heading: str | None = None

    # Running headers of every page this chunk spans, in page order. A list
    # rather than a single value: a chunk crossing a chapter boundary would
    # otherwise silently keep only the first page's chapter.
    chapters: list[str] = Field(default_factory=list)

    # Breadcrumb of enclosing headings, outermost first, built from MinerU's
    # `text_level` hierarchy. Example from the BUMA standard:
    #   ["2. PENJELASAN PARAMETER", "2.1. Production Parameter", "2.1.1. Production"]
    heading_path: list[str] = Field(default_factory=list)

    # Flags for cheap filtering downstream
    has_formula: bool = False
    is_tabular: bool = False

    # Trace back to the source: item indices in MinerU's content_list.json
    source_items: list[int] = Field(default_factory=list)

    # Position on the page of the first source item, as MinerU reports it.
    # Carried through for a curation UI that highlights where on the page a
    # definition came from. Nothing in the pipeline reasons about it — ordering
    # uses item sequence, never coordinates.
    bbox: list[int] | None = None

    # Figures, charts, diagrams and tables this chunk contains or points at.
    # Added 0.4.0 — see the `Asset` docstring for why images stopped being chunks.
    #
    # A prose chunk that holds a figure carries it here AND carries the figure's
    # `![](asset://<asset_id>)` reference inline in `text` at its reading-order
    # position, so the model reads the reference exactly where the document put it.
    assets: list[Asset] = Field(default_factory=list)

    # chunk_ids of OTHER chunks that point at this one. Populated on table chunks,
    # which stay standalone (their linearised text earns its place as evidence —
    # 74.6% of the Komatsu manual's top-3 evidence is tabular) while still being
    # reachable from the prose that introduces them.
    referenced_by: list[str] = Field(default_factory=list)

    # ⚠️ DEPRECATED 0.4.0 — superseded by `assets`. Still populated for one
    # version so a 0.3.x reader degrades rather than breaking; it holds the bare
    # parser-local filenames, which are NOT unique across a document (see
    # `Asset.asset_id`). Read `assets` instead.
    images: list[str] = Field(default_factory=list)

    # Source markup, kept verbatim beside the rendered prose in `text`.
    #
    # `text` carries a readable rendering because the term filter is an NER
    # model reading prose: MinerU writes formulas character-spaced
    # ("P u r c h a s i n g ~ c o s t s") and tables as HTML, and neither
    # produces a single mention. Measured on the same document and gold set,
    # only the parse differing: raw markup in `text` scored recall 0.7561 against
    # 0.8537 for plain text; rendering it back recovered 0.8293.
    #
    # The markup is not discarded — the formula branch needs exactly this form.
    latex: list[str] = Field(default_factory=list)
    table_html: str | None = None

    # ⚠️ MODEL-WRITTEN, NOT FROM THE DOCUMENT. MinerU's `vlm` backend (and
    # `hybrid` at effort=high) describes images and charts; that description
    # arrives in the item's `content` field.
    #
    # It is kept OUT of `text` deliberately. `text` is the haystack the
    # extraction guardrail searches to prove a quote is verbatim; putting
    # generated prose there would let a hallucinated sentence pass the very
    # check that exists to catch hallucination. A caption printed in the
    # document is verbatim and belongs in `text`; a description invented about
    # the picture does not.
    #
    # Any consumer using this must treat it as a model claim about the image,
    # never as evidence — it cannot be span-checked against anything.
    generated_description: str | None = None


class ParsedDocument(BaseModel):
    """The artifact itself — one parsed document, self-describing.

    The chunk list alone is not enough to hand across the seam: an artifact that
    travels to the extraction half must carry its own provenance. Without
    `parser_name` / `parser_version` / `parser_backend`, a MinerU upgrade and a
    prompt change are indistinguishable when extraction results shift.

    `content_hash` is the hash of the SOURCE FILE, so re-parsing the same
    document is detectable and a changed document forces a new artifact version.

    **Metadata is declared before `chunks`, and that ordering is deliberate.**
    Pydantic serialises in declaration order, so it decides what a person sees on
    opening the file. With `chunks` first, checking which parser or which
    normaliser produced an artifact means scrolling past every chunk to the end —
    thousands of lines in a real document. That is not hypothetical fussiness:
    verifying `parser_config` and `normaliser_version` is now a routine step when
    an artifact is handed between the two halves, and burying them is how a stale
    artifact went unnoticed for a day.

    JSON objects are unordered by specification and every consumer reads by key,
    so this changes nothing for code — only for the person reading it.
    """

    doc_id: str

    # The document's name as a HUMAN reads it. Today this is the file stem and
    # therefore equal to `doc_id` — the point is that it is a separate field, so
    # the two can diverge without anything having to change here.
    #
    # `doc_id` is on its way to becoming an identifier rather than a label (see
    # the parsing README): every `term_id` / `formula_id` / `rule_id` downstream
    # is a digest of it, so it must stop depending on what someone typed into a
    # Save dialog. Once it does, `Provenance` — which is what a reviewing expert
    # actually reads — would otherwise show them a hash instead of the standard's
    # name. This field is what travels alongside it.
    #
    # Added ahead of that switch because it is needed whichever identifier wins,
    # and adding it now costs nothing and breaks nothing.
    source_title: str | None = None

    # Identity of the source
    source_path: str
    content_hash: str                    # sha256 of the source file
    n_pages: int

    # Which parser produced this, and how
    parser_name: str = "mineru"
    parser_version: str | None = None    # e.g. "3.4.4"
    parser_backend: str | None = None    # "pipeline" | "vlm" | "hybrid", as MinerU recorded it
    parser_config: str | None = None     # fingerprint of the settings that affect output

    # Fingerprint of the code that turned MinerU's output into these chunks —
    # `normalize.py`, `render.py`, `checks.py`, `contracts.py`. OPTIONAL and
    # purely informational: nothing downstream has to read it, and an artifact
    # written before this field existed simply carries null.
    #
    # It closes a real gap. `parser_config` fingerprints MinerU's settings, and
    # `schema_version` versions the artifact's SHAPE — neither moves when a
    # normalisation bug is fixed, because the shape is unchanged. So a stale
    # artifact and a correct one declared identical metadata, and the only
    # observable difference was the character count. That cost a day on
    # 2026-08-26 (see `normalize.normaliser_version`).
    normaliser_version: str | None = None

    # Version of THIS ARTIFACT for this document — bumped when the document is
    # re-parsed (new source content, new parser version, or changed settings).
    # Distinct from `schema_version`, which versions the contract itself.
    version: int = 1

    schema_version: str = SCHEMA_VERSION
    created_at: str | None = None

    # Where the untouched MinerU output for this document lives
    raw_output_dir: str | None = None

    # LAST on purpose — see the class docstring. Everything above answers "what
    # is this file?", and a reader should not have to scroll past the payload to
    # find out.
    chunks: list[Chunk]


# --- Seam for the downstream stages (declared, not yet used) ---
# Written here so the handoff shape is visible from the start. The extraction
# half owns the final form of these two — see `KNOWLEDGE_PIPELINE_TODO.md` §5.


class Mention(BaseModel):
    """One occurrence of a term inside a chunk."""

    term: str
    chunk_id: str
    page_idx: int


class TermRecord(BaseModel):
    """Extracted knowledge for one term (one cluster of mentions)."""

    term: str
    full_name: str | None = None

    # What the document literally says, before any normalisation — e.g.
    # "Physical of Availability (PA)" where `full_name` may read "Physical
    # Availability". Span-checked against `Chunk.text` like every other field,
    # so the discrepancy surfaces to the expert instead of being quietly fixed.
    source_wording: str | None = None

    definition: str | None = None
    formula_latex: str | None = None
    subdomain_tags: list[str] = Field(default_factory=list)
    mention_count: int = 0
    provenance: dict = Field(default_factory=dict)
    extraction_status: str = "ok"