"""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")) 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, ) # 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 def complete(self) -> bool: 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.""" head = text[:4000] number = DOC_NUMBER_RE.search(head) revision = REVISION_RE.search(head) effective = DATE_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, ) 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)'}") 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, ) # 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 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): log.info( "keeping a %d-token procedure section whole - it contains a step sequence", approx_tokens(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 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, ) ) 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, embedding) VALUES (%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())