yau-plant-assistant/api/chunking.py

"""Sectioning, chunking and header extraction — MIRRORED FROM ingest/ingest.py.

WHY THIS FILE IS A COPY, WHICH IS NORMALLY THE WRONG ANSWER.

Documents reach `doc_chunks` by two routes now: `ai-ingest` at a terminal, and
the upload screen in this API. **They must chunk identically.** If they do not,
the same document ingested by the two paths produces different chunks,
different embeddings and different retrieval behaviour — and the difference
would show up as an assistant that answers a question correctly or not
depending on who loaded the document, which is close to undiagnosable.

The functions below are therefore mirrored from `ingest/ingest.py` with the
logic unchanged. They are not imported from it because `ingest.py` lives in a
different image, with a different and heavier dependency set.
A shared module needs both images built from a common context, which is a
deployment change to a live host and not one to make in the same stroke as a
new feature.

**If you change chunking, change it in BOTH files.** `api/tests/test_chunking.py`
locks the behaviour that matters — the refusal to split a numbered step
sequence — so a drift in that rule fails the suite. It cannot detect drift in
`ingest.py`, which is the residual risk and is recorded in
BUILD-AI-CONTAINERS.md S14.

Mirrored: MARKDOWN_HEADING, STEP_RE, the header regexes, approx_tokens,
split_sections, chunk_section, link_equipment, extract_header, parse_date.
NOT mirrored: parse_document, and everything that touches the disk.
"""

from __future__ import annotations

import os
import re
from dataclasses import dataclass
from datetime import date, datetime

CHUNK_TOKEN_TARGET = int(os.environ.get("CHUNK_TOKEN_TARGET", "800"))

# The hard ceiling no chunk may exceed, whatever the input looks like.
#
# Set well under text-embedding-3-small's 8191-token input limit, because
# approx_tokens() is a length/4 estimate and under-counts dense technical text.
# A chunk over the model's limit does not degrade - the embeddings call fails
# and the document cannot be published at all.
MAX_CHUNK_TOKENS = int(os.environ.get("MAX_CHUNK_TOKENS", "6000"))

MARKDOWN_HEADING = re.compile(r"^(#{1,6})\s+(.*\S)\s*$")

DOC_NUMBER_RE = re.compile(r"\b(WRPS-[A-Z]{2,4}-\d{3,4})\b")
REVISION_RE = re.compile(
    r"\b(?:rev(?:ision)?|issue)[\s.:]*([A-Z0-9]{1,4})\b", re.IGNORECASE
)
DATE_RE = re.compile(
    r"\b(?:effective|issued|approved)[\s\w]{0,12}?[:\s]\s*"
    r"(\d{1,2}[/-]\d{1,2}[/-]\d{2,4}|\d{4}-\d{2}-\d{2}|"
    r"\d{1,2}\s+\w+\s+\d{4})\b",
    re.IGNORECASE,
)
# [ 	], NOT \s: \s includes the newline, so the old form matched a bare
# "Title" line, consumed the line break and captured whatever was on the NEXT
# line. A real cover page flattens to a label column then a value column -
# Title / Revision / Date - and every chunk of WRPS-CTL-001 was stored with
# doc_title = "Revision". Same trap AUTHORISING_ROLE_RE was fixed for.
TITLE_RE = re.compile(r"^[ 	]*title[ 	:]+(.+\S)[ 	]*$",
                      re.IGNORECASE | re.MULTILINE)
AUTHORISING_ROLE_RE = re.compile(
    r"^[ \t]*authoris(?:ing|ed)[ \t]*(?:role|by)?[ \t]*:[ \t]*(.+\S)[ \t]*$",
    re.IGNORECASE | re.MULTILINE,
)

# A numbered step. Used to refuse to split, not to parse the procedure.
STEP_RE = re.compile(r"^\s*(?:\d+\.|\(\d+\)|step\s+\d+)", re.IGNORECASE | re.MULTILINE)


@dataclass(frozen=True)
class Header:
    doc_number: str | None
    revision: str | None
    effective_date: date | None
    title: str | None
    authorising_role: str | None


def parse_date(text: str) -> date | None:
    for fmt in ("%d/%m/%Y", "%d-%m-%Y", "%Y-%m-%d", "%d %B %Y", "%d %b %Y", "%d/%m/%y"):
        try:
            return datetime.strptime(text.strip(), fmt).date()
        except ValueError:
            continue
    return None


def extract_header(text: str) -> Header:
    """Pull document identity from the first page. ALWAYS confirmed by a human.

    Everything here is a proposal shown on the review screen with the field
    already filled in. It is not authority. A wrong revision on a procedure is
    a safety issue, and a regex is not a person.
    """
    head = text[:4000]
    number = DOC_NUMBER_RE.search(head)
    revision = REVISION_RE.search(head)
    effective = DATE_RE.search(head)
    title = TITLE_RE.search(head)
    role = AUTHORISING_ROLE_RE.search(head)
    return Header(
        doc_number=number.group(1) if number else None,
        revision=revision.group(1) if revision else None,
        effective_date=parse_date(effective.group(1)) if effective else None,
        title=title.group(1).strip() if title else None,
        authorising_role=role.group(1).strip() if role else None,
    )


def approx_tokens(text: str) -> int:
    """Rough, and deliberately so - it decides when to split, and the rule that
    matters is the one that refuses to."""
    return len(text) // 4


def split_sections(text: str) -> list[tuple[int, str, str]]:
    """Markdown -> [(page, section_title, section_text)].

    Mirrors ingest.parse_markdown, which takes a Path; this takes the text,
    because by the time it gets here the upload has already been converted in
    memory.

    Page is always 1. A converted document has no pages that survive the
    conversion, so a citation to one carries a section title and no page
    number. That is honest. Do not invent page numbers to make citations look
    uniform.
    """
    sections: list[tuple[int, str, list[str]]] = []
    current_title = "(untitled)"
    buffer: list[str] = []

    for line in text.splitlines():
        heading = MARKDOWN_HEADING.match(line)
        if heading:
            if buffer:
                sections.append((1, current_title, buffer))
            current_title = heading.group(2).strip()
            buffer = []
        elif line.strip():
            buffer.append(line.rstrip())

    if buffer:
        sections.append((1, current_title, buffer))

    return [(page, title, "\n".join(body)) for page, title, body in sections]


def _split_on_words(line: str, limit: int) -> list[str]:
    """Split one over-long line on word boundaries.

    The bottom of the ladder. A PDF page can extract as a SINGLE line with no
    newline anywhere in it, and at that point there is no structure left to
    respect - only the limit, which is not negotiable because exceeding the
    embedding model's input makes the document unpublishable.
    """
    pieces: list[str] = []
    buffer: list[str] = []
    for word in line.split(" "):
        candidate = " ".join(buffer + [word])
        if buffer and approx_tokens(candidate) > limit:
            pieces.append(" ".join(buffer))
            buffer = [word]
        else:
            buffer.append(word)
    if buffer:
        pieces.append(" ".join(buffer))
    # A single "word" longer than the limit is not language - it is a base64
    # blob or a table rendered without spaces. Slice it rather than emit it.
    bounded: list[str] = []
    for piece in pieces:
        while approx_tokens(piece) > limit:
            bounded.append(piece[: limit * 4])
            piece = piece[limit * 4 :]
        if piece:
            bounded.append(piece)
    return bounded


def _split_on_lines(text: str, limit: int) -> list[str]:
    """Last-resort split, on line boundaries, honouring `limit`.

    Used when paragraph splitting could not get a chunk under the ceiling -
    text with no blank lines in it at all, which is exactly what flat PDF
    extraction produces. Falls through to word boundaries for a single line
    that is itself over the limit.
    """
    chunks: list[str] = []
    buffer: list[str] = []
    for line in text.splitlines() or [text]:
        if approx_tokens(line) > limit:
            if buffer:
                chunks.append("\n".join(buffer))
                buffer = []
            chunks.extend(_split_on_words(line, limit))
            continue
        candidate = "\n".join(buffer + [line])
        if buffer and approx_tokens(candidate) > limit:
            chunks.append("\n".join(buffer))
            buffer = [line]
        else:
            buffer.append(line)
    if buffer:
        chunks.append("\n".join(buffer))
    return chunks


def chunk_section(text: str, doc_type: str) -> list[str]:
    """Split a section, unless splitting it would break a step sequence.

    For procedures the rule still holds: a section containing numbered steps is
    emitted whole rather than split. An oversized chunk costs tokens; half a
    procedure costs more than that.

    THE CEILING, added 2026-08-28. The rule above used to be unbounded, and
    that was safe only while sections arrived pre-split by Docling and were
    therefore small. They no longer always do: flat PDF extraction can hand
    this function an entire document as one section, and an unbounded refusal
    then produces one chunk for the whole document - one embedding vector for
    eight pages, useless retrieval, and a citation reading "(untitled)". Worse,
    a long enough document exceeds the embedding model's input limit and the
    publish fails outright.

    So the refusal is now bounded by MAX_CHUNK_TOKENS. Below it, a procedure
    section stays whole exactly as before. Above it, splitting is the lesser
    harm - an embedding call that fails protects nobody. The gap between
    CHUNK_TOKEN_TARGET and MAX_CHUNK_TOKENS is deliberately wide so that a real
    step sequence, which is the case the rule exists for, is never near it.
    """
    if approx_tokens(text) <= CHUNK_TOKEN_TARGET:
        return [text]

    if doc_type == "procedure" and STEP_RE.search(text):
        if approx_tokens(text) <= MAX_CHUNK_TOKENS:
            return [text]
        # Past the ceiling. Fall through and split - and split on lines, since
        # a blob this shape usually has no paragraph breaks to use.
        return _split_on_lines(text, CHUNK_TOKEN_TARGET)

    chunks: list[str] = []
    buffer: list[str] = []
    for paragraph in text.split("\n\n"):
        candidate = "\n\n".join(buffer + [paragraph])
        if buffer and approx_tokens(candidate) > CHUNK_TOKEN_TARGET:
            chunks.append("\n\n".join(buffer))
            buffer = [paragraph]
        else:
            buffer.append(paragraph)
    if buffer:
        chunks.append("\n\n".join(buffer))

    # A single paragraph can still be over the ceiling - text with no blank
    # lines is one paragraph however long it is. Nothing above this point can
    # fix that, so enforce it here rather than trusting the input.
    bounded: list[str] = []
    for chunk in chunks:
        if approx_tokens(chunk) > MAX_CHUNK_TOKENS:
            bounded.extend(_split_on_lines(chunk, CHUNK_TOKEN_TARGET))
        else:
            bounded.append(chunk)
    return bounded


def link_equipment(text: str, equipment_ids: list[str]) -> str | None:
    """Tie a chunk to equipment when it is unambiguously about one thing.

    Two different units mentioned means no link, not a guess - a chunk linked
    to the wrong pump is worse than one linked to nothing, because retrieval
    filtering will then hide it from the pump it actually describes.
    """
    found = {eid for eid in equipment_ids if eid.lower() in text.lower()}
    return found.pop() if len(found) == 1 else None