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

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,
)
TITLE_RE = re.compile(r"^\s*title[\s:]+(.+\S)\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 chunk_section(text: str, doc_type: str) -> list[str]:
    """Split a section, unless splitting it would break a step sequence.

    For procedures the rule is absolute: a section containing numbered steps is
    emitted whole, however long it is. An oversized chunk costs tokens. Half a
    procedure costs more than that.
    """
    if approx_tokens(text) <= CHUNK_TOKEN_TARGET:
        return [text]

    if doc_type == "procedure" and STEP_RE.search(text):
        return [text]

    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))
    return chunks


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