diff --git a/ingest/ingest.py b/ingest/ingest.py index 3054246..5851d95 100644 --- a/ingest/ingest.py +++ b/ingest/ingest.py @@ -146,13 +146,60 @@ def confirm_header(path: Path, header: Header, assume_yes: bool) -> Header: ) +# 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]]: - """Docling -> [(page, section_title, section_text)]. + """[(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. + 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)) @@ -259,7 +306,14 @@ def superseded_state(conn: psycopg.Connection, source_file: str) -> bool | None: return row[0] if row else None -def ingest_file(path: Path, conn: psycopg.Connection, client: AzureOpenAI, *, assume_yes: bool) -> int: +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: @@ -297,7 +351,14 @@ def ingest_file(path: Path, conn: psycopg.Connection, client: AzureOpenAI, *, as ) ) - vectors = embed_all([r[9] for r in records], client, os.environ["EMBED_DEPLOYMENT"]) + 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. @@ -312,7 +373,10 @@ def ingest_file(path: Path, conn: psycopg.Connection, client: AzureOpenAI, *, as embedding) VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s) """, - [record + (str(vector),) for record, vector in zip(records, vectors)], + [ + 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), @@ -458,13 +522,33 @@ def main() -> int: "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() - 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"], - ) + 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) @@ -503,7 +587,12 @@ def main() -> int: else: parser.error("give --all or --file") - total = sum(ingest_file(p, conn, client, assume_yes=args.assume_yes) for p in paths) + 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