34d2ccc576
Build spec and host brief carried in from C:\Claude and WRPS/02-env; the plant model (equipment, tags, alarm bitmask, enums, unit conversions) is derived from WRPS/04-plc/register-map.csv, WRPS/05-scada/modbus/scada-points.csv and WRPS-CTL-003. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
361 lines
13 KiB
Python
361 lines
13 KiB
Python
"""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.
|
|
"""
|
|
|
|
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,
|
|
)
|
|
|
|
|
|
def parse_document(path: Path) -> list[tuple[int, str, str]]:
|
|
"""Docling -> [(page, section_title, section_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.
|
|
"""
|
|
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 ingest_file(path: Path, conn: psycopg.Connection, client: AzureOpenAI, *, assume_yes: bool) -> 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))
|
|
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, False,
|
|
link_equipment(chunk, equipment_ids), page, title, chunk,
|
|
)
|
|
)
|
|
|
|
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),) 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 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",
|
|
)
|
|
args = parser.parse_args()
|
|
|
|
dsn = (
|
|
f"postgresql://{os.environ['PGUSER']}:{os.environ['PGPASSWORD']}"
|
|
f"@{os.environ['PGHOST']}:{os.environ.get('PGPORT','5432')}"
|
|
f"/{os.environ['PGDATABASE']}"
|
|
)
|
|
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(dsn, application_name="ai-ingest") as conn:
|
|
if args.supersede:
|
|
count = mark_superseded(conn, *args.supersede)
|
|
conn.commit()
|
|
log.info("marked %d chunks superseded", 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"}
|
|
)
|
|
else:
|
|
parser.error("give --all or --file")
|
|
|
|
total = sum(ingest_file(p, conn, client, assume_yes=args.assume_yes) for p in paths)
|
|
|
|
log.info("done: %d chunks from %d files", total, len(paths))
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main())
|