dcfb411c3b
ProcedureIdentity requires a title and an authorising role, and neither was stored anywhere. The answer writer was asked for both, read them off whatever chunk retrieval happened to return, and returned "" whenever the header chunk was not among them. They belong in the row for the same reason doc_number and revision do: they are facts about the controlled document, established once when a human confirms the header, not something to re-derive per question from whatever text was retrieved. Denormalised onto every chunk exactly as the existing header fields are - ingest replaces every chunk of a source_file in one transaction, so they cannot drift within a document. complete() deliberately still requires only doc_number, revision and effective_date. 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. controlled_copy_location is NOT in the schema. It is a site fact, identical on every row, and the one field where an invented value sends a person to a place that does not exist. It is CONTROLLED_COPY_LOCATION in api.env, defaulting to a string that names who to ask. The authorising-role pattern requires the colon: without it the lazy gap swallowed the field name and captured "role: Station Maintenance Supervisor" as the value, which the first run caught. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
632 lines
24 KiB
Python
632 lines
24 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.
|
|
|
|
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,
|
|
)
|
|
# 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 "".
|
|
TITLE_RE = re.compile(r"^\s*title[\s:]+(.+\S)\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."""
|
|
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 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,
|
|
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())
|