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

ChunkKind = Literal["text", "table", "chart", "equation"]

SCHEMA_VERSION = "0.2.0"


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

    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 MinerU reports
    # them, with no conversion anywhere in the pipeline. Page numbers eventually
    # reach a human review queue, and an off-by-one there is invisible until an
    # expert opens the wrong page. Converting to 1-based is the UI's job, done
    # once at display time — never here, so the artifact always matches the raw
    # MinerU output kept alongside it.
    page_idx: int
    page_idxs: list[int] = Field(default_factory=list)

    # 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

    # Non-text attachments (formula images, table/chart crops)
    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


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.
    """

    doc_id: str
    chunks: list[Chunk]

    # 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

    # 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


# --- 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"