yau-plant-assistant/ingest/ingest.py

"""Document ingestion: Docling parse -> chunk -> embed -> pg-ai.

Run on demand, not as a service:

    docker compose -f ~/ai-compose.yml run --rm ai-ingest --all
    docker compose -f ~/ai-compose.yml run --rm ai-ingest --file procedures/WRPS-OPS-014.pdf

Documents live on /datadisk/ai-docs, mounted read-only at /docs. They are NOT
in Git - the repo's docs/ directory is a gitignored placeholder.

    /docs/procedures/        doc_type = procedure
    /docs/manuals/           doc_type = manual
    /docs/rationalisation/   doc_type = rationalisation
    /docs/design/            doc_type = design

FOUR RULES, in descending order of how badly it goes if you break them:

1. A WRONG REVISION ON A PROCEDURE IS A SAFETY ISSUE, not a data quality one.
   doc_number, revision and effective_date are extracted from the header and
   then CONFIRMED BY A HUMAN before the chunks are committed. --assume-yes
   exists for re-ingesting already-confirmed files and nothing else.

2. NEVER SPLIT A NUMBERED STEP SEQUENCE ACROSS CHUNKS. If a section exceeds the
   token target, keep it whole. Half a step sequence retrieved on its own is
   how a partial procedure reaches somebody.

3. doc_type COMES FROM THE FOLDER, never from the model, never from the file
   name. A manual filed under procedures/ is a filing error to fix on disk.

4. RE-RUNS REPLACE, NEVER DUPLICATE. Chunks for a source_file are deleted and
   reinserted in one transaction.

5. RE-INGESTING NEVER RESURRECTS A WITHDRAWN DOCUMENT. A superseded document
   comes back superseded, and --all skips it entirely. Replacing chunks used to
   reset the flag to FALSE, so one bulk re-run quietly made every withdrawn
   revision citable again - including the old revision of a procedure.
"""

from __future__ import annotations

import argparse
import logging
import os
import re
import sys
from dataclasses import dataclass
from datetime import date, datetime
from pathlib import Path

import psycopg
from openai import AzureOpenAI

logging.basicConfig(level=logging.INFO, format="%(levelname)s %(message)s")
log = logging.getLogger("ingest")

DOCS_ROOT = Path(os.environ.get("AI_DOCS_ROOT", "/docs"))
CHUNK_TOKEN_TARGET = int(os.environ.get("CHUNK_TOKEN_TARGET", "800"))
# The hard ceiling no chunk may exceed. Mirrors api/chunking.py - the two must
# stay in step or the same document chunks differently depending on which path
# loaded it. Set well under text-embedding-3-small's 8191-token input limit,
# because approx_tokens() is a length/4 estimate and under-counts dense 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"))

DOC_TYPE_BY_FOLDER = {
    "procedures": "procedure",
    "manuals": "manual",
    "rationalisation": "rationalisation",
    "design": "design",
}

# WRPS document numbering: WRPS-CTL-001, WRPS-PRO-001, WRPS-OPS-014, WRPS-DRG-001.
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,
)
# Header fields that are not safety-critical but ARE facts about the document:
# storing them stops the model being asked to read them off whatever chunk
# retrieval happened to return, which is how it ended up returning "".
# [ 	], 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: X", "Authorised by: X", "Authorising: X". The colon is
# required: without it the lazy gap swallowed the field NAME and captured
# "role: Station Maintenance Supervisor" as the value.
AUTHORISING_ROLE_RE = re.compile(
    r"^[ 	]*authoris(?:ing|ed)[ 	]*(?:role|by)?[ 	]*:[ 	]*(.+\S)[ 	]*$",
    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
class Header:
    doc_number: str | None
    revision: str | None
    effective_date: date | None
    title: str | None = None
    authorising_role: str | None = None

    def complete(self) -> bool:
        """The three fields a wrong value in is a SAFETY issue.

        Deliberately not title or authorising_role: a missing title makes an
        answer less useful, a wrong revision sends somebody to the wrong
        document. --assume-yes must keep refusing on the second and tolerate
        the first.
        """
        return all((self.doc_number, self.revision, self.effective_date))


def doc_type_for(path: Path) -> str:
    try:
        folder = path.relative_to(DOCS_ROOT).parts[0]
    except ValueError:
        folder = path.parent.name
    if folder not in DOC_TYPE_BY_FOLDER:
        raise SystemExit(
            f"{path}: folder {folder!r} is not one of {sorted(DOC_TYPE_BY_FOLDER)}. "
            "doc_type comes from the folder - move the file, do not override this."
        )
    return DOC_TYPE_BY_FOLDER[folder]


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 confirm_header(path: Path, header: Header, assume_yes: bool) -> Header:
    """Ask a person. A wrong revision on a procedure is a safety issue."""
    print(f"\n{path}")
    print(f"  doc_number     : {header.doc_number or '(not found)'}")
    print(f"  revision       : {header.revision or '(not found)'}")
    print(f"  effective_date : {header.effective_date or '(not found)'}")
    print(f"  title          : {header.title or '(not found)'}")
    print(f"  authorising    : {header.authorising_role or '(not found)'}")

    if assume_yes:
        if not header.complete():
            raise SystemExit(
                f"{path}: --assume-yes but the header is incomplete. Confirm it "
                "by hand - this is the field where a mistake is a safety issue."
            )
        return header

    if input("  Correct? [y/N] ").strip().lower() == "y":
        return header
    return Header(
        doc_number=input("  doc_number     : ").strip() or header.doc_number,
        revision=input("  revision       : ").strip() or header.revision,
        effective_date=parse_date(input("  effective_date (YYYY-MM-DD): ").strip())
        or header.effective_date,
        title=input("  title          : ").strip() or header.title,
        authorising_role=input("  authorising    : ").strip() or header.authorising_role,
    )


# Formats whose structure is already explicit in the bytes. Running a document
# layout model over a file that literally contains "## 2. Prerequisites" buys
# nothing, and it is the difference between an image with torch in it and one
# without.
PLAIN_TEXT_SUFFIXES = {".md", ".markdown", ".txt"}

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


def parse_markdown(path: Path) -> list[tuple[int, str, str]]:
    """Markdown/plain text -> [(page, section_title, section_text)].

    Sections split on ATX headings, which is the same section-boundary rule
    Docling applies to a PDF - so chunk_section() sees the same shape either
    way and the "never split a step sequence" rule still holds.

    Page is always 1: a Markdown file has no pages. A citation to it carries a
    section title and no meaningful page number, which is honest. Do not invent
    page numbers to make citations look uniform.
    """
    text = path.read_text(encoding="utf-8")
    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 parse_document(path: Path) -> list[tuple[int, str, str]]:
    """[(page, section_title, section_text)]. Docling, unless it is plain text.

    Docling gives structure, which is what makes section-boundary chunking
    possible. A plain text extractor would force splitting on token count, and
    token-count splitting is what cuts step sequences in half. That reasoning
    applies to PDFs and Word documents; for Markdown the structure is already
    in the file, so parse_markdown does the same job without importing a
    machine learning stack.
    """
    if path.suffix.lower() in PLAIN_TEXT_SUFFIXES:
        return parse_markdown(path)

    from docling.document_converter import DocumentConverter

    result = DocumentConverter().convert(str(path))
    document = result.document

    sections: list[tuple[int, str, list[str]]] = []
    current_title = "(untitled)"
    current_page = 1
    buffer: list[str] = []

    for item, _level in document.iterate_items():
        text = getattr(item, "text", "") or ""
        if not text.strip():
            continue
        page = getattr(getattr(item, "prov", [None])[0], "page_no", current_page) or current_page
        label = str(getattr(item, "label", "")).lower()

        if "header" in label or "title" in label or "section" in label:
            if buffer:
                sections.append((current_page, current_title, buffer))
            current_title = text.strip()
            current_page = page
            buffer = []
        else:
            buffer.append(text)
            current_page = page

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

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


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_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


def embed_all(texts: list[str], client: AzureOpenAI, model: str) -> list[list[float]]:
    vectors: list[list[float]] = []
    for i in range(0, len(texts), 64):        # batch, to keep the call count sane
        batch = texts[i : i + 64]
        response = client.embeddings.create(model=model, input=batch)
        vectors.extend(item.embedding for item in response.data)
    return vectors


def superseded_state(conn: psycopg.Connection, source_file: str) -> bool | None:
    """Is this file already ingested, and was it withdrawn? None = not ingested.

    bool_or, not bool_and: if any chunk of the file is superseded the document
    is treated as superseded. The conservative direction is the one that keeps a
    withdrawn procedure out of an answer.
    """
    with conn.cursor() as cur:
        cur.execute(
            "SELECT bool_or(superseded) FROM doc_chunks WHERE source_file = %s",
            (source_file,),
        )
        row = cur.fetchone()
    return row[0] if row else None


def ingest_file(
    path: Path,
    conn: psycopg.Connection,
    client: AzureOpenAI | None,
    *,
    assume_yes: bool,
    no_embed: bool = False,
) -> int:
    doc_type = doc_type_for(path)
    sections = parse_document(path)
    if not sections:
        log.warning("%s: nothing extracted - check the file", path)
        return 0

    full_text = "\n".join(body for _, _, body in sections)
    header = confirm_header(path, extract_header(full_text), assume_yes)

    with conn.cursor() as cur:
        cur.execute("SELECT equipment_id FROM equipment")
        equipment_ids = [row[0] for row in cur.fetchall()]

    records: list[tuple] = []
    source_file = str(path.relative_to(DOCS_ROOT))

    # Withdrawal survives a re-ingest. Replacing the chunks must not silently
    # give this document its citability back - if it was superseded before, it
    # is superseded after, and saying so out loud is the point.
    superseded = superseded_state(conn, source_file) or False
    if superseded:
        log.warning(
            "%s is currently SUPERSEDED - re-ingesting it as superseded. It "
            "will not be cited. Use --restore %s %s to bring it back.",
            source_file, header.doc_number, header.revision,
        )

    for page, title, body in sections:
        for chunk in chunk_section(body, doc_type):
            records.append(
                (
                    source_file, doc_type, header.doc_number, header.revision,
                    header.effective_date, superseded,
                    link_equipment(chunk, equipment_ids), page, title, chunk,
                    header.title, header.authorising_role,
                )
            )

    if no_embed:
        # NULL, not a zero vector. A zero vector is a POINT in the space and
        # ranks against real queries - it would surface as a plausible hit for
        # anything. NULL sorts last and returns no similarity at all, which is
        # the honest representation of "this chunk has not been embedded".
        vectors = [None] * len(records)
    else:
        vectors = embed_all([r[9] for r in records], client, os.environ["EMBED_DEPLOYMENT"])

    # Replace, never duplicate - both statements in one transaction, so a
    # failure halfway does not leave the document half-ingested.
    with conn.transaction():
        with conn.cursor() as cur:
            cur.execute("DELETE FROM doc_chunks WHERE source_file = %s", (source_file,))
            cur.executemany(
                """
                INSERT INTO doc_chunks
                    (source_file, doc_type, doc_number, revision, effective_date,
                     superseded, equipment_id, page, section_title, chunk_text,
                     doc_title, authorising_role, embedding)
                VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s)
                """,
                [
                    record + (str(vector) if vector is not None else None,)
                    for record, vector in zip(records, vectors)
                ],
            )

    log.info("%s: %d chunks (%s rev %s)", source_file, len(records),
             header.doc_number, header.revision)
    return len(records)


def mark_superseded(conn: psycopg.Connection, doc_number: str, keep_revision: str) -> int:
    """Withdraw every revision of a document except the current one.

    Retrieval filters superseded = FALSE, so this is how an old revision stops
    being citable. Run it whenever a new revision is ingested - the ingest does
    not infer it, because inferring which revision is current from a header is
    exactly the judgement that needs a person.
    """
    with conn.cursor() as cur:
        cur.execute(
            "UPDATE doc_chunks SET superseded = TRUE"
            " WHERE doc_number = %s AND revision <> %s AND superseded = FALSE",
            (doc_number, keep_revision),
        )
        return cur.rowcount


def connection_params() -> dict[str, str | int]:
    """Connect as the role that can WRITE, not the one the API uses.

    PGUSER in ~/ai/api.env is agent_ro - SELECT and nothing else, deliberately,
    because it is the role the answer path runs as. ai-ingest reads the same
    env file, so it inherited that role and could not insert a chunk. Ingestion
    connects as ingest_rw (db/003_roles.sql) via INGEST_DB_USER.

    The fallback to PGUSER exists for a local shell where only the one pair is
    set. It is not a way to run ingestion as agent_ro - require_write_access()
    below refuses that, whichever variable it came from.

    Keyword parameters rather than a URL: a password containing @ or / breaks a
    DSN string silently, and these are generated passwords.
    """
    return {
        "host": os.environ["PGHOST"],
        "port": int(os.environ.get("PGPORT", "5432")),
        "dbname": os.environ["PGDATABASE"],
        "user": os.environ.get("INGEST_DB_USER") or os.environ["PGUSER"],
        "password": os.environ.get("INGEST_DB_PASSWORD") or os.environ["PGPASSWORD"],
    }


def require_write_access(conn: psycopg.Connection) -> None:
    """Refuse early, before anything is parsed, embedded or paid for.

    Without this the run does the whole job - Docling parse, header
    confirmation typed by a person, an embeddings call that is billed - and
    then fails on the INSERT with a permission error that names no cause. The
    check is one round trip and it fails with the fix in it.
    """
    try:
        with conn.cursor() as cur:
            cur.execute(
                "SELECT current_user,"
                "       has_table_privilege('doc_chunks', 'INSERT'),"
                "       has_table_privilege('doc_chunks', 'UPDATE'),"
                "       has_table_privilege('doc_chunks', 'DELETE')"
            )
            user, may_insert, may_update, may_delete = cur.fetchone()
    except psycopg.errors.UndefinedTable as missing:
        raise SystemExit(
            "doc_chunks does not exist in this database. Apply db/001_schema.sql "
            "and db/003_roles.sql first - see README.md, Phase 1."
        ) from missing

    if may_insert and may_update and may_delete:
        log.info("connected as %s", user)
        return

    raise SystemExit(
        f"connected to pg-ai as {user!r}, which cannot write doc_chunks. "
        "Ingestion must connect as ingest_rw. Set INGEST_DB_USER and "
        "INGEST_DB_PASSWORD in ~/ai/api.env - PGUSER there is agent_ro, which "
        "is SELECT-only on purpose and must stay that way. "
        "If the role does not exist yet, apply db/003_roles.sql."
    )


def restore(conn: psycopg.Connection, doc_number: str, revision: str) -> int:
    """Bring a withdrawn revision back. The counterpart of --supersede.

    Refused while another revision of the same document is live. Restoring
    rev 3 next to rev 4 puts two revisions of one procedure in front of an
    operator, which is the failure the superseded filter exists to prevent -
    and it is a likelier mistake than it sounds, because the person restoring
    is usually looking at the old revision, not the new one.
    """
    with conn.cursor() as cur:
        cur.execute(
            "SELECT DISTINCT revision FROM doc_chunks"
            " WHERE doc_number = %s AND revision <> %s AND superseded = FALSE",
            (doc_number, revision),
        )
        live = [row[0] for row in cur.fetchall()]
        if live:
            raise SystemExit(
                f"{doc_number} revision {', '.join(live)} is live. Restoring "
                f"revision {revision} would put two revisions of one document in "
                "front of an operator. Supersede the other one first, if that is "
                "really what you mean."
            )

        cur.execute(
            "UPDATE doc_chunks SET superseded = FALSE"
            " WHERE doc_number = %s AND revision = %s AND superseded = TRUE",
            (doc_number, revision),
        )
        return cur.rowcount


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--all", action="store_true", help="ingest every document")
    parser.add_argument("--file", help="one file, relative to the docs root")
    parser.add_argument(
        "--assume-yes",
        action="store_true",
        help="skip header confirmation - only for files already confirmed once",
    )
    parser.add_argument(
        "--supersede",
        nargs=2,
        metavar=("DOC_NUMBER", "KEEP_REVISION"),
        help="mark every other revision of a document superseded",
    )
    parser.add_argument(
        "--restore",
        nargs=2,
        metavar=("DOC_NUMBER", "REVISION"),
        help="bring a withdrawn revision back; refused if another revision is live",
    )
    parser.add_argument(
        "--include-superseded",
        action="store_true",
        help=(
            "with --all, do not skip withdrawn documents. They are still "
            "re-ingested AS withdrawn - this only spends the embedding call"
        ),
    )
    parser.add_argument(
        "--no-embed",
        action="store_true",
        help=(
            "insert chunks with a NULL embedding and make no Azure OpenAI call. "
            "For NO_LLM_STUB demos before an OpenAI account exists. These chunks "
            "are INVISIBLE to vector search and findable only lexically - "
            "re-ingest properly once embeddings are available"
        ),
    )
    args = parser.parse_args()

    if args.no_embed:
        # Loud, because a database half full of unembeddable chunks looks
        # exactly like working retrieval right up until it silently returns
        # nothing for the question that matters.
        log.warning("--no-embed: chunks will have NO EMBEDDING and CANNOT be found")
        log.warning("by semantic search. Re-ingest every document without this")
        log.warning("flag once EMBED_DEPLOYMENT is configured. Clear them first:")
        log.warning("  DELETE FROM doc_chunks WHERE embedding IS NULL;")
        client = None
    else:
        client = AzureOpenAI(
            azure_endpoint=os.environ["AZURE_OPENAI_ENDPOINT"],
            api_key=os.environ["AZURE_OPENAI_API_KEY"],
            api_version=os.environ["AZURE_OPENAI_API_VERSION"],
        )

    with psycopg.connect(**connection_params(), application_name="ai-ingest") as conn:
        require_write_access(conn)
        if args.supersede:
            count = mark_superseded(conn, *args.supersede)
            conn.commit()
            log.info("marked %d chunks superseded", count)
            return 0

        if args.restore:
            count = restore(conn, *args.restore)
            conn.commit()
            log.info("restored %d chunks", count)
            return 0

        if args.file:
            paths = [DOCS_ROOT / args.file]
        elif args.all:
            paths = sorted(
                p
                for folder in DOC_TYPE_BY_FOLDER
                for p in (DOCS_ROOT / folder).glob("**/*")
                if p.is_file() and p.suffix.lower() in {".pdf", ".docx", ".md", ".txt"}
            )
            # A withdrawn document is still sitting in the tree - nothing moves
            # it. Skipping it keeps a bulk re-run from spending an embeddings
            # call on a document that will not be cited either way.
            if not args.include_superseded:
                keep = []
                for p in paths:
                    if superseded_state(conn, str(p.relative_to(DOCS_ROOT))):
                        log.info("skipping %s - superseded", p.relative_to(DOCS_ROOT))
                    else:
                        keep.append(p)
                paths = keep
        else:
            parser.error("give --all or --file")

        total = sum(
            ingest_file(
                p, conn, client, assume_yes=args.assume_yes, no_embed=args.no_embed
            )
            for p in paths
        )

    log.info("done: %d chunks from %d files", total, len(paths))
    return 0


if __name__ == "__main__":
    sys.exit(main())